Add network context manager

Author: mloesch

core/testcontainers/core/network.py

@@ -0,0 +1,44 @@
+#
+#    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.
+from typing import Optional
+
+from testcontainers.core.docker_client import DockerClient
+
+
+class Network(object):
+    """
+    Network context manager to conveniently connect containers.
+    """
+
+    def __init__(self, name, docker_client_kw: Optional[dict] = None, **kwargs) -> None:
+        self.name = name
+        self._docker = DockerClient(**(docker_client_kw or {}))
+        self._kwargs = kwargs
+
+    def remove(self) -> None:
+        self._network.remove()
+
+    def __enter__(self) -> 'Network':
+        self._network = self._docker.client.networks.create(self.name, **self._kwargs)
+        self.id = self._network.id
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.remove()
+
+    def __del__(self) -> None:
+        if self._network is not None:
+            try:
+                self.remove()
+            except:  # noqa: E722
+                pass

core/tests/test_network.py

@@ -0,0 +1,27 @@
+from testcontainers.core.container import DockerContainer
+from testcontainers.core.docker_client import DockerClient
+from testcontainers.core.network import Network
+
+
+def test_network_gets_created_and_cleaned_up():
+    with Network("test-network") as network:
+        docker = DockerClient()
+        networks_list = docker.client.networks.list("test-network")
+        assert networks_list[0].name == "test-network"
+        assert networks_list[0].id == network.id
+    assert not docker.client.networks.list("test-network")
+
+
+def test_containers_can_communicate_over_network():
+    with Network("network") as network:
+        with DockerContainer("nginx:alpine-slim").with_name(
+                "alpine1").with_kwargs(network=network.name) as alpine1:
+            with DockerContainer("nginx:alpine-slim").with_name(
+                    "alpine2").with_kwargs(network=network.name) as alpine2:
+                status, output = alpine1.exec("ping -c 1 alpine2")
+                assert status == 0
+                assert "64 bytes" in str(output)
+
+                status, output = alpine2.exec("ping -c 1 alpine1")
+                assert status == 0
+                assert "64 bytes" in str(output)