Add tests for container routes (#475)

* Add tests for container routes

* Clean up async tests

* Improve run tests

* Load an actual pipeline in tests rather than mocking

* Add more tests

* Update black version

* remove legacy environments commands and run black

* Bump version

.github/workflows/pull_request.yml

@@ -7,7 +7,7 @@ jobs:
       - uses: actions/checkout@v3
       - uses: psf/black@stable
         with:
-          version: 23.12.0
+          version: 24.4.2
 
   flake8-lint:
     runs-on: ubuntu-latest
@@ -61,4 +61,4 @@ jobs:
           poetry install --no-interaction
       - name: test
         run: |
-          poetry run pytest
+          PYTHONASYNCIODEBUG=1 poetry run pytest

examples/audio-to-text/whisper-large/new_pipeline.py

@@ -24,8 +24,6 @@ class ModelKwargs(InputSchema):
 
 @entity
 class WhisperModel:
-    def __init__(self):
-        ...
 
     @pipe(on_startup=True, run_once=True)
     def load(self):

examples/image-to-image/t2i-adapter-sketch/new_pipeline.py

@@ -108,8 +108,6 @@ class ModelKwargs(InputSchema):
 
 @entity
 class DiffusionWithAdapter:
-    def __init__(self) -> None:
-        ...
 
     def apply_style(
         self, style_name: str, positive: str, negative: str = ""

examples/text-to-audio/musicgen-large/new_pipeline.py

@@ -4,8 +4,6 @@
 
 @entity
 class MusicgenModel:
-    def __init__(self):
-        ...
 
     @pipe(on_startup=True, run_once=True)
     def load(self):

pipeline/console/cluster.py

@@ -52,7 +52,7 @@ def use_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
 
 
 def remove_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
-    ...
+    pass
 
 
 def get_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:

pipeline/console/commands.py

@@ -2,7 +2,6 @@
 
 from pipeline.console import cluster, container, logs
 from pipeline.console.targets import (
-    environments,
     files,
     pipelines,
     pointers,
@@ -23,7 +22,6 @@ def create_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
         dest="target",
     )
 
-    environments.create_parser(create_sub_parser)
     pointers.create_parser(create_sub_parser)
     resources.create_parser(create_sub_parser)
     files.create_parser(create_sub_parser)
@@ -42,7 +40,6 @@ def edit_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
         dest="target",
     )
 
-    environments.edit_parser(edit_sub_parser)
     pipelines.edit_parser(edit_sub_parser)
     pointers.edit_parser(edit_sub_parser)
     files.edit_parser(edit_sub_parser)
@@ -61,7 +58,6 @@ def get_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
         dest="target",
     )
 
-    environments.get_parser(get_sub_parser)
     pipelines.get_parser(get_sub_parser)
     pointers.get_parser(get_sub_parser)
     resources.get_parser(get_sub_parser)
@@ -82,7 +78,6 @@ def delete_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
         dest="target",
     )
 
-    environments.delete_parser(delete_sub_parser)
     pipelines.delete_parser(delete_sub_parser)
     pointers.delete_parser(delete_sub_parser)
     resources.delete_parser(delete_sub_parser)

pipeline/console/targets/files.py

@@ -41,7 +41,7 @@ def create_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
 
 
 def edit_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
-    ...
+    pass
 
 
 def get_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
@@ -73,7 +73,7 @@ def get_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
 
 
 def delete_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
-    ...
+    pass
 
 
 def _create_file(args: Namespace) -> None:

pipeline/console/targets/resources.py

@@ -10,7 +10,7 @@
 
 
 def create_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
-    ...
+    pass
 
 
 def get_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
@@ -21,7 +21,7 @@ def get_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
 
 
 def delete_parser(command_parser: "_SubParsersAction[ArgumentParser]") -> None:
-    ...
+    pass
 
 
 # Functions

pipeline/container/startup.py

@@ -2,6 +2,7 @@
 import os
 import traceback
 import uuid
+from contextlib import asynccontextmanager
 
 import pkg_resources
 from fastapi import FastAPI, Request
@@ -20,24 +21,13 @@
 
 
 def create_app() -> FastAPI:
-    app = FastAPI(
-        title="pipeline-container",
-    )
+    app = FastAPI(title="pipeline-container", lifespan=lifespan)
 
     setup_logging()
 
     setup_oapi(app)
     setup_middlewares(app)
 
-    app.state.execution_queue = asyncio.Queue()
-    app.state.manager = Manager(
-        pipeline_path=os.environ.get(
-            "PIPELINE_PATH",
-            "",
-        )
-    )
-    asyncio.create_task(execution_handler(app.state.execution_queue, app.state.manager))
-
     app.include_router(router)
     app.include_router(status_router)
     static_dir = pkg_resources.resource_filename(
@@ -53,6 +43,22 @@ def create_app() -> FastAPI:
     return app
 
 
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.state.execution_queue = asyncio.Queue()
+    app.state.manager = Manager(
+        pipeline_path=os.environ.get(
+            "PIPELINE_PATH",
+            "",
+        )
+    )
+    task = asyncio.create_task(
+        execution_handler(app.state.execution_queue, app.state.manager)
+    )
+    yield
+    task.cancel()
+
+
 def setup_middlewares(app: FastAPI) -> None:
     @app.middleware("http")
     async def _(request: Request, call_next):

pipeline/objects/pipeline.py

@@ -6,11 +6,6 @@ class Pipeline:
     _current_pipeline: Graph
     _pipeline_context_active: bool = False
 
-    def __init__(
-        self,
-    ):
-        ...
-
     def __enter__(self):
         Pipeline._pipeline_context_active = True
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pipeline-ai"
-version = "2.4.6"
+version = "2.4.7"
 description = "Pipelines for machine learning workloads."
 authors = [
   "Paul Hetherington <ph@mystic.ai>",
@@ -32,7 +32,7 @@ fastapi = "^0.105.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.3"
-black = "^23.12.0"
+black = "^24.4.2"
 pre-commit = "^3.6.0"
 flake8 = "^6.1.0"
 isort = "^5.13.2"

tests/container/fixtures/adder_pipeline.py

@@ -0,0 +1,21 @@
+import logging
+
+from pipeline import Pipeline, Variable, pipe
+
+logger = logging.getLogger(__name__)
+
+
+@pipe
+def add(first: int, second: int) -> int:
+    if first < 0 or second < 0:
+        raise ValueError("I can only sum positive integers")
+    return first + second
+
+
+with Pipeline() as builder:
+    first = Variable(int)
+    second = Variable(int)
+    result = add(first, second)
+    builder.output(result)
+
+my_pipeline = builder.get_pipeline()

tests/container/routes/conftest.py

@@ -0,0 +1,16 @@
+import pytest
+from fastapi.testclient import TestClient
+
+from pipeline.container.startup import create_app
+
+
+@pytest.fixture
+async def app():
+    app = create_app()
+    return app
+
+
+@pytest.fixture
+async def client(app):
+    with TestClient(app) as client:
+        yield client