Use docker networks for integration tests

This is possible now we don't have file upload server.

Closes replicate#262
Closes replicate#263

Signed-off-by: Ben Firshman <[email protected]>

Author: bfirsh

test-integration/test_integration/conftest.py

@@ -1,25 +1,43 @@
 import subprocess
 import pytest
+import redis
 
-from .util import random_string, find_free_port, docker_run, get_local_ip, wait_for_port
+from .util import random_string, find_free_port, docker_run, wait_for_port
 
 
 @pytest.fixture
-def redis_port():
-    container_name = "cog-test-redis-" + random_string(10)
+def redis_port(docker_network):
+    """Start a redis server inside the Docker network.
+    Inside the network, it is available at redis:6379
+    Outside the network, it is available at localhost:redis_port
+    """
     port = find_free_port()
     with docker_run(
         "redis",
-        name=container_name,
+        net_alias="redis",
+        network=docker_network,
         publish=[{"host": port, "container": 6379}],
         detach=True,
     ):
-        wait_for_port(get_local_ip(), port)
+        wait_for_port("localhost", port)
         yield port
 
 
+@pytest.fixture
+def redis_client(redis_port):
+    yield redis.Redis("localhost", redis_port)
+
+
 @pytest.fixture
 def docker_image():
     image = "cog-test-" + random_string(10)
     yield image
     subprocess.run(["docker", "rmi", "-f", image], check=False)
+
+
+@pytest.fixture
+def docker_network():
+    name = "cog-test-" + random_string(10)
+    subprocess.run(["docker", "network", "create", name])
+    yield name
+    subprocess.run(["docker", "network", "rm", name])

test-integration/test_integration/test_redis_queue.py

@@ -1,34 +1,25 @@
 import json
 from pathlib import Path
 
-import redis
 import subprocess
 
-from .util import (
-    docker_run,
-    get_local_ip,
-    random_string,
-)
+from .util import docker_run, random_string
 
 
-def test_queue_worker_yielding(docker_image, redis_port):
+def test_queue_worker_yielding(docker_image, docker_network, redis_client):
     project_dir = Path(__file__).parent / "fixtures/yielding-project"
     subprocess.run(["cog", "build", "-t", docker_image], check=True, cwd=project_dir)
 
-    local_ip = get_local_ip()
-    redis_host = local_ip
-
-    redis_client = redis.Redis(host=redis_host, port=redis_port, db=0)
-
     with docker_run(
         image=docker_image,
         interactive=True,
+        network=docker_network,
         command=[
             "python",
             "-m",
             "cog.server.redis_queue",
-            redis_host,
-            str(redis_port),
+            "redis",
+            "6379",
             "predict-queue",
             "",
             "test-worker",
@@ -69,24 +60,20 @@ def test_queue_worker_yielding(docker_image, redis_port):
         assert response == None
 
 
-def test_queue_worker_error(docker_image, redis_port, request):
+def test_queue_worker_error(docker_image, docker_network, redis_client):
     project_dir = Path(__file__).parent / "fixtures/failing-project"
     subprocess.run(["cog", "build", "-t", docker_image], check=True, cwd=project_dir)
 
-    local_ip = get_local_ip()
-    redis_host = local_ip
-
-    redis_client = redis.Redis(host=redis_host, port=redis_port, db=0)
-
     with docker_run(
         image=docker_image,
         interactive=True,
+        network=docker_network,
         command=[
             "python",
             "-m",
             "cog.server.redis_queue",
-            redis_host,
-            str(redis_port),
+            "redis",
+            "6379",
             "predict-queue",
             "",
             "test-worker",
@@ -121,24 +108,20 @@ def test_queue_worker_error(docker_image, redis_port, request):
         assert response == None
 
 
-def test_queue_worker_logging(docker_image, redis_port, request):
+def test_queue_worker_logging(docker_image, docker_network, redis_client):
     project_dir = Path(__file__).parent / "fixtures/logging-project"
     subprocess.run(["cog", "build", "-t", docker_image], check=True, cwd=project_dir)
 
-    local_ip = get_local_ip()
-    redis_host = local_ip
-
-    redis_client = redis.Redis(host=redis_host, port=redis_port, db=0)
-
     with docker_run(
         image=docker_image,
         interactive=True,
+        network=docker_network,
         command=[
             "python",
             "-m",
             "cog.server.redis_queue",
-            redis_host,
-            str(redis_port),
+            "redis",
+            "6379",
             "predict-queue",
             "",
             "test-worker",
@@ -189,24 +172,20 @@ def test_queue_worker_logging(docker_image, redis_port, request):
         ]
 
 
-def test_queue_worker_timeout(docker_image, redis_port, request):
+def test_queue_worker_timeout(docker_image, docker_network, redis_client):
     project_dir = Path(__file__).parent / "fixtures/timeout-project"
     subprocess.run(["cog", "build", "-t", docker_image], check=True, cwd=project_dir)
 
-    local_ip = get_local_ip()
-    redis_host = local_ip
-
-    redis_client = redis.Redis(host=redis_host, port=redis_port, db=0)
-
     with docker_run(
         image=docker_image,
         interactive=True,
+        network=docker_network,
         command=[
             "python",
             "-m",
             "cog.server.redis_queue",
-            redis_host,
-            str(redis_port),
+            "redis",
+            "6379",
             "predict-queue",
             "",
             "test-worker",
@@ -258,24 +237,20 @@ def test_queue_worker_timeout(docker_image, redis_port, request):
         assert response == {"status": "failed", "error": "Prediction timed out"}
 
 
-def test_queue_worker_yielding_timeout(docker_image, redis_port, request):
+def test_queue_worker_yielding_timeout(docker_image, docker_network, redis_client):
     project_dir = Path(__file__).parent / "fixtures/yielding-timeout-project"
     subprocess.run(["cog", "build", "-t", docker_image], check=True, cwd=project_dir)
 
-    local_ip = get_local_ip()
-    redis_host = local_ip
-
-    redis_client = redis.Redis(host=redis_host, port=redis_port, db=0)
-
     with docker_run(
         image=docker_image,
         interactive=True,
+        network=docker_network,
         command=[
             "python",
             "-m",
             "cog.server.redis_queue",
-            redis_host,
-            str(redis_port),
+            "redis",
+            "6379",
             "predict-queue",
             "",
             "test-worker",

test-integration/test_integration/util.py

@@ -38,6 +38,8 @@ def docker_run(
     name=None,
     detach=False,
     interactive=False,
+    network=None,
+    net_alias=None,
     publish: Optional[List[dict]] = None,
     command: Optional[List[str]] = None,
     env: Optional[dict] = None,
@@ -46,6 +48,7 @@ def docker_run(
         name = random_string(10)
 
     cmd = ["docker", "run", "--name", name]
+
     if publish is not None:
         for port_binding in publish:
             host_port = port_binding["host"]
@@ -58,6 +61,10 @@ def docker_run(
         cmd += ["--detach"]
     if interactive:
         cmd += ["-i"]
+    if network is not None:
+        cmd.extend(["--network", network])
+    if net_alias is not None:
+        cmd.extend(["--net-alias", net_alias])
     cmd += [image]
     if command:
         cmd += command
@@ -66,21 +73,3 @@ def docker_run(
         yield
     finally:
         subprocess.Popen(["docker", "rm", "--force", name]).wait()
-
-
-def get_local_ip():
-    return socket.gethostbyname(socket.gethostname())
-
-
-def get_bridge_ip():
-    """Return the IP address of the docker bridge network"""
-    cmd = [
-        "docker",
-        "network",
-        "inspect",
-        "bridge",
-        "--format='{{ json (index .IPAM.Config 0).Gateway }}'",
-    ]
-    out, _ = subprocess.Popen(cmd, stdout=subprocess.PIPE).communicate()
-
-    return out.decode().strip().replace('"', "").replace("'", "")