Skip to content

Commit

Permalink
feat(network): Add network context manager (#367)
Browse files Browse the repository at this point in the history
This PR adds a `Network` helper class that allows to create networks and
connect containers programmatically.

The networks are context-managed resources like containers created via
`DockerContainer`.

Please also see tests for a usage example :)

---------

Co-authored-by: Kevin Wittek <kiview@users.noreply.github.com>
Co-authored-by: Max Pfeiffer <max@maxpfeiffer.ch>
Co-authored-by: Jakob Beckmann <32326425+f4z3r@users.noreply.github.com>
Co-authored-by: David Ankin <daveankin@gmail.com>
Co-authored-by: Carli* Freudenberg <carli.freudenberg@energymeteo.de>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Vemund Santi <vemund@santi.no>
  • Loading branch information
8 people committed Apr 16, 2024
1 parent 29b5179 commit 11964de
Show file tree
Hide file tree
Showing 3 changed files with 97 additions and 0 deletions.
13 changes: 13 additions & 0 deletions core/testcontainers/core/container.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from testcontainers.core.docker_client import DockerClient
from testcontainers.core.exceptions import ContainerStartException
from testcontainers.core.labels import LABEL_SESSION_ID, SESSION_ID
from testcontainers.core.network import Network
from testcontainers.core.utils import inside_container, is_arm, setup_logger
from testcontainers.core.waiting_utils import wait_container_is_ready, wait_for_logs

Expand Down Expand Up @@ -46,6 +47,8 @@ def __init__(
self._container = None
self._command = None
self._name = None
self._network: Optional[Network] = None
self._network_aliases: Optional[list[str]] = None
self._kwargs = kwargs

def with_env(self, key: str, value: str) -> Self:
Expand All @@ -61,6 +64,14 @@ def with_exposed_ports(self, *ports: int) -> Self:
self.ports[port] = None
return self

def with_network(self, network: Network) -> Self:
self._network = network
return self

def with_network_aliases(self, *aliases) -> Self:
self._network_aliases = aliases
return self

def with_kwargs(self, **kwargs) -> Self:
self._kwargs = kwargs
return self
Expand All @@ -87,6 +98,8 @@ def start(self) -> Self:
**self._kwargs,
)
logger.info("Container started: %s", self._container.short_id)
if self._network:
self._network.connect(self._container.id, self._network_aliases)
return self

def stop(self, force=True, delete_volume=True) -> None:
Expand Down
41 changes: 41 additions & 0 deletions core/testcontainers/core/network.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
import uuid
from typing import Optional

from testcontainers.core.docker_client import DockerClient


class Network:
"""
Network context manager for programmatically connecting containers.
"""

def __init__(self, docker_client_kw: Optional[dict] = None, docker_network_kw: Optional[dict] = None) -> None:
self.name = str(uuid.uuid4())
self._docker = DockerClient(**(docker_client_kw or {}))
self._docker_network_kw = docker_network_kw or {}

def connect(self, container_id: str, network_aliases: Optional[list] = None):
self._network.connect(container_id, aliases=network_aliases)

def remove(self) -> None:
self._network.remove()

def __enter__(self) -> "Network":
self._network = self._docker.client.networks.create(self.name, **self._docker_network_kw)
self.id = self._network.id
return self

def __exit__(self, exc_type, exc_val, exc_tb) -> None:
self.remove()
43 changes: 43 additions & 0 deletions core/tests/test_network.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
from testcontainers.core.container import DockerContainer
from testcontainers.core.docker_client import DockerClient
from testcontainers.core.network import Network

NGINX_ALPINE_SLIM_IMAGE = "nginx:1.25.4-alpine-slim"


def test_network_gets_created_and_cleaned_up():
with Network() as network:
docker = DockerClient()
networks_list = docker.client.networks.list(network.name)
assert networks_list[0].name == network.name
assert networks_list[0].id == network.id
assert not docker.client.networks.list(network.name)


def test_containers_can_communicate_over_network():
with Network() as network:
with (
DockerContainer(NGINX_ALPINE_SLIM_IMAGE)
.with_name("alpine1")
.with_network_aliases("alpine1-alias-1", "alpine1-alias-2")
.with_network(network) as alpine1
):
with (
DockerContainer(NGINX_ALPINE_SLIM_IMAGE)
.with_name("alpine2")
.with_network_aliases("alpine2-alias-1", "alpine2-alias-2")
.with_network(network) as alpine2
):
assert_can_ping(alpine1, "alpine2")
assert_can_ping(alpine1, "alpine2-alias-1")
assert_can_ping(alpine1, "alpine2-alias-2")

assert_can_ping(alpine2, "alpine1")
assert_can_ping(alpine2, "alpine1-alias-1")
assert_can_ping(alpine2, "alpine1-alias-2")


def assert_can_ping(container: DockerContainer, remote_name: str):
status, output = container.exec("ping -c 1 %s" % remote_name)
assert status == 0
assert "64 bytes" in str(output)

0 comments on commit 11964de

Please sign in to comment.