fix auto deploy

lm-cyber · Apr 11, 2024 · a0bc902 · a0bc902
1 parent 65ec91a
commit a0bc902
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/app/main.py b/app/main.py
@@ -6,7 +6,7 @@
 from fastapi import FastAPI, Request
 from fastapi.logger import logger
 from fastapi.middleware.cors import CORSMiddleware
-
+from tests import api_test
 from model import Model
 from predict import predict
 from config import CONFIG
@@ -16,7 +16,6 @@
     InferenceOutput,
     ErrorResponse,
 )
-from tests import test_get_plot, test_post_date
 
 
 @asynccontextmanager
@@ -110,8 +109,7 @@ def bash(command):
 
 
 ################ Test views #################
-app.post("/api/v1/test_post_date")(test_post_date)
-app.get("/api/v1/test_get_plot")(test_get_plot)
+app.include_router(api_test, prefix="/api/v1", tags=["tests"])
 #############################################
 
 if __name__ == "__main__":

diff --git a/app/tests.py b/app/tests.py
@@ -1,12 +1,18 @@
 import random
 
 from schema import TestInput
+from fastapi import APIRouter
 
 
+api_test = APIRouter(tags=["test"])
+
+
+@api_test.get("/test_get_plot")
 def test_post_date(data: TestInput):
     return data
 
 
+@api_test.post("/test_post_date")
 def test_get_plot() -> dict[str, list[int]]:
     arr_len = 100
     x = [i for i in range(arr_len)]