diff --git a/core/testcontainers/core/container.py b/core/testcontainers/core/container.py index efa06734..d2605490 100644 --- a/core/testcontainers/core/container.py +++ b/core/testcontainers/core/container.py @@ -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 @@ -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: @@ -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 @@ -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: diff --git a/core/testcontainers/core/network.py b/core/testcontainers/core/network.py new file mode 100644 index 00000000..9903d071 --- /dev/null +++ b/core/testcontainers/core/network.py @@ -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() diff --git a/core/tests/test_network.py b/core/tests/test_network.py new file mode 100644 index 00000000..4b0764d4 --- /dev/null +++ b/core/tests/test_network.py @@ -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)