From 17c5c1133323bdd125c66422f391dd261618b306 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lucas=20Serv=C3=A9n=20Mar=C3=ADn?= <lserven@gmail.com>
Date: Thu, 18 Apr 2024 22:27:38 -0700
Subject: [PATCH] feat: add test to exercise Nomad scheduler (#182)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Lucas Servén Marín <lserven@gmail.com>
---
 projects/fal/tests/test_apps.py      | 42 ++++++++++++++++++++++++++--
 projects/fal/tests/test_stability.py | 19 +++++++++++--
 2 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/projects/fal/tests/test_apps.py b/projects/fal/tests/test_apps.py
index c42ddf87..2945eaf4 100644
--- a/projects/fal/tests/test_apps.py
+++ b/projects/fal/tests/test_apps.py
@@ -8,7 +8,8 @@
 from fastapi import WebSocket
 from httpx import HTTPStatusError
 from openapi_fal_rest.api.applications import app_metadata
-from pydantic import BaseModel, __version__ as pydantic_version
+from pydantic import BaseModel
+from pydantic import __version__ as pydantic_version
 
 import fal
 import fal.api as api
@@ -47,6 +48,23 @@ def addition_app(input: Input) -> Output:
     return Output(result=input.lhs + input.rhs)
 
 
+@fal.function(
+    keep_alive=60,
+    machine_type="S",
+    serve=True,
+    max_concurrency=1,
+    requirements=[f"pydantic=={pydantic_version}"],
+    _scheduler="nomad",
+)
+def nomad_addition_app(input: Input) -> Output:
+    print("starting...")
+    for _ in range(input.wait_time):
+        print("sleeping...")
+        time.sleep(1)
+
+    return Output(result=input.lhs + input.rhs)
+
+
 @fal.function(
     keep_alive=300,
     requirements=["fastapi", "uvicorn", "pydantic==1.10.12"],
@@ -185,6 +203,20 @@ def test_app():
     yield f"{user_id}/{app_revision}"
 
 
+@pytest.fixture(scope="module")
+def test_nomad_app():
+    # Create a temporary app, register it, and return the ID of it.
+
+    from fal.cli import _get_user_id
+
+    app_revision = nomad_addition_app.host.register(
+        func=nomad_addition_app.func,
+        options=nomad_addition_app.options,
+    )
+    user_id = _get_user_id()
+    yield f"{user_id}/{app_revision}"
+
+
 @pytest.fixture(scope="module")
 def test_fastapi_app():
     # Create a temporary app, register it, and return the ID of it.
@@ -230,13 +262,19 @@ def test_realtime_app():
     yield f"{user_id}/{app_revision}"
 
 
-def test_app_client(test_app: str):
+def test_app_client(test_app: str, test_nomad_app: str):
     response = apps.run(test_app, arguments={"lhs": 1, "rhs": 2})
     assert response["result"] == 3
 
     response = apps.run(test_app, arguments={"lhs": 2, "rhs": 3, "wait_time": 1})
     assert response["result"] == 5
 
+    response = apps.run(test_nomad_app, arguments={"lhs": 1, "rhs": 2})
+    assert response["result"] == 3
+
+    response = apps.run(test_nomad_app, arguments={"lhs": 2, "rhs": 3, "wait_time": 1})
+    assert response["result"] == 5
+
 
 def test_app_client_old_format(test_app: str):
     assert test_app.count("/") == 1, "Test app should be in new format"
diff --git a/projects/fal/tests/test_stability.py b/projects/fal/tests/test_stability.py
index c31faa97..f63a5cc2 100644
--- a/projects/fal/tests/test_stability.py
+++ b/projects/fal/tests/test_stability.py
@@ -14,7 +14,6 @@
 
 
 def test_missing_dependencies_nested_server_error(isolated_client):
-
     @isolated_client()
     def test1():
         return "hello"
@@ -44,6 +43,20 @@ def mult(a, b):
     assert mult(5, 2) == 10
 
 
+def test_regular_function_on_nomad(isolated_client):
+    @isolated_client(_scheduler="nomad")
+    def regular_function():
+        return 42
+
+    assert regular_function() == 42
+
+    @isolated_client(_scheduler="nomad")
+    def mult(a, b):
+        return a * b
+
+    assert mult(5, 2) == 10
+
+
 def test_function_pipelining(isolated_client):
     @isolated_client("virtualenv")
     def regular_function():
@@ -433,7 +446,9 @@ class MathQuery(BaseModel):
     class MathResult(BaseModel):
         result: int = Field(description="The result of the operation")
 
-    @isolated_client("virtualenv", serve=True, requirements=[f"pydantic=={pydantic_version}"])
+    @isolated_client(
+        "virtualenv", serve=True, requirements=[f"pydantic=={pydantic_version}"]
+    )
     def add(query: MathQuery) -> MathResult:
         return MathResult(result=query.x + query.y)