Test queue worker and refactor end to end tests

Signed-off-by: andreasjansson <[email protected]>

Author: andreasjansson

end-to-end-test/end_to_end_test/__init__.py

@@ -0,0 +1,101 @@
+import os
+import subprocess
+import tempfile
+from waiting import wait
+import requests
+import pytest
+
+from .util import random_string, find_free_port, docker_run
+
+
+@pytest.fixture
+def cog_server_port():
+    old_cwd = os.getcwd()
+    with tempfile.TemporaryDirectory() as cog_dir:
+        os.chdir(cog_dir)
+        port = str(find_free_port())
+        server_proc = subprocess.Popen(["cog", "server", "--port", port])
+        resp = wait(
+            lambda: requests.get("http://localhost:" + port + "/ping"),
+            timeout_seconds=60,
+            expected_exceptions=(requests.exceptions.ConnectionError,),
+        )
+        assert resp.text == "pong"
+
+        yield port
+
+    os.chdir(old_cwd)
+    server_proc.kill()
+
+
+@pytest.fixture
+def project_dir(tmpdir_factory):
+    tmpdir = tmpdir_factory.mktemp("project")
+    with open(tmpdir / "infer.py", "w") as f:
+        f.write(
+            """
+import time
+import tempfile
+from pathlib import Path
+import cog
+
+class Model(cog.Model):
+    def setup(self):
+        self.foo = "foo"
+
+    @cog.input("text", type=str)
+    @cog.input("path", type=Path)
+    @cog.input("output_file", type=bool, default=False)
+    def run(self, text, path, output_file):
+        time.sleep(1)
+        with open(path) as f:
+            output = self.foo + text + f.read()
+        if output_file:
+            tmp = tempfile.NamedTemporaryFile(suffix=".txt")
+            tmp.close()
+            tmp_path = Path(tmp.name)
+            with tmp_path.open("w") as f:
+                f.write(output)
+                return tmp_path
+        return output
+        """
+        )
+    with open(tmpdir / "cog.yaml", "w") as f:
+        cog_yaml = """
+name: andreas/hello-world
+model: infer.py:Model
+examples:
+  - input:
+      text: "foo"
+      path: "@myfile.txt"
+    output: "foofoobaz"
+  - input:
+      text: "bar"
+      path: "@myfile.txt"
+    output: "foobarbaz"
+  - input:
+      text: "qux"
+      path: "@myfile.txt"
+environment:
+  architectures:
+    - cpu
+        """
+        f.write(cog_yaml)
+
+    with open(tmpdir / "myfile.txt", "w") as f:
+        f.write("baz")
+
+    return tmpdir
+
+
+@pytest.fixture
+def redis_port():
+    container_name = "cog-test-redis-" + random_string(10)
+    port = find_free_port()
+    with docker_run(
+        "redis",
+        name=container_name,
+        publish=[{"host": port, "container": 6379}],
+        detach=True,
+    ):
+        yield port

end-to-end-test/end_to_end_test/conftest.py

@@ -0,0 +1,150 @@
+import json
+import redis
+from contextlib import contextmanager
+import multiprocessing
+from flask import Flask, request, jsonify, Response
+
+from .util import (
+    random_string,
+    set_model_url,
+    show_version,
+    push_with_log,
+    find_free_port,
+    docker_run,
+    get_local_ip,
+    wait_for_port,
+)
+
+
+def test_queue_worker(cog_server_port, project_dir, redis_port, tmpdir_factory):
+    user = random_string(10)
+    model_name = random_string(10)
+    model_url = f"http://localhost:{cog_server_port}/{user}/{model_name}"
+
+    set_model_url(model_url, project_dir)
+    version_id = push_with_log(project_dir)
+    version_data = show_version(model_url, version_id)
+
+    input_queue = multiprocessing.Queue()
+    output_queue = multiprocessing.Queue()
+    controller_port = find_free_port()
+    local_ip = get_local_ip()
+    upload_url = f"http://{local_ip}:{controller_port}/upload"
+    redis_host = local_ip
+    worker_name = "test-worker"
+    infer_queue_name = "infer-queue"
+    response_queue_name = "response-queue"
+
+    wait_for_port(redis_host, redis_port)
+
+    redis_client = redis.Redis(host=redis_host, port=redis_port, db=0)
+
+    with queue_controller(input_queue, output_queue, controller_port), docker_run(
+        image=version_data["images"][0]["uri"],
+        interactive=True,
+        command=[
+            "cog-redis-queue-worker",
+            redis_host,
+            str(redis_port),
+            infer_queue_name,
+            upload_url,
+            worker_name,
+        ],
+    ):
+        redis_client.xgroup_create(
+            mkstream=True, groupname=infer_queue_name, name=infer_queue_name, id="$"
+        )
+
+        infer_id = random_string(10)
+        redis_client.xadd(
+            name=infer_queue_name,
+            fields={
+                "value": json.dumps(
+                    {
+                        "id": infer_id,
+                        "inputs": {
+                            "text": {"value": "bar"},
+                            "path": {
+                                "file": {
+                                    "name": "myinput.txt",
+                                    "url": f"http://{local_ip}:{controller_port}/download",
+                                }
+                            },
+                        },
+                        "response_queue": response_queue_name,
+                    }
+                ),
+            },
+        )
+        input_queue.put("test")
+        response = json.loads(redis_client.brpop(response_queue_name)[1])["value"]
+        assert response == "foobartest"
+
+        infer_id = random_string(10)
+        redis_client.xadd(
+            name=infer_queue_name,
+            fields={
+                "value": json.dumps(
+                    {
+                        "id": infer_id,
+                        "inputs": {
+                            "text": {"value": "bar"},
+                            "output_file": {"value": "true"},
+                            "path": {
+                                "file": {
+                                    "name": "myinput.txt",
+                                    "url": f"http://{local_ip}:{controller_port}/download",
+                                }
+                            },
+                        },
+                        "response_queue": response_queue_name,
+                    }
+                ),
+            },
+        )
+        input_queue.put("test")
+        response_contents = output_queue.get()
+        response = json.loads(redis_client.brpop(response_queue_name)[1])["file"]
+        assert response_contents.decode() == "foobartest"
+        assert response["url"] == "uploaded.txt"
+
+
+@contextmanager
+def queue_controller(input_queue, output_queue, controller_port):
+    controller = QueueController(input_queue, output_queue, controller_port)
+    controller.start()
+    yield controller
+    controller.kill()
+
+
+class QueueController(multiprocessing.Process):
+    def __init__(self, input_queue, output_queue, port):
+        super().__init__()
+        self.input_queue = input_queue
+        self.output_queue = output_queue
+        self.port = port
+
+    def run(self):
+        app = Flask("test-queue-controller")
+
+        @app.route("/", methods=["GET"])
+        def handle_index():
+            return "OK"
+
+        @app.route("/upload", methods=["PUT"])
+        def handle_upload():
+            f = request.files["file"]
+            contents = f.read()
+            self.output_queue.put(contents)
+            return jsonify({"url": "uploaded.txt"})
+
+        @app.route("/download", methods=["GET"])
+        def handle_download():
+            contents = self.input_queue.get()
+            return Response(
+                contents,
+                mimetype="text/plain",
+                headers={"Content-disposition": "attachment; filename=myinput.txt"},
+            )
+
+        app.run(host="0.0.0.0", port=self.port, debug=False)