From 3a656dcc995f20c60a6541758127a2bfbd7cdc8f Mon Sep 17 00:00:00 2001
From: Pao Sheng <paooap.oappao@gmail.com>
Date: Wed, 5 Jun 2024 15:13:04 +0800
Subject: [PATCH] chore(wren-ai-service): move the development purpose API to
 the development router (#381)

---
 wren-ai-service/src/__main__.py            | 22 +++++-----------------
 wren-ai-service/src/web/development.py     | 20 ++++++++++++++++++++
 wren-ai-service/tests/locust/locustfile.py |  2 +-
 3 files changed, 26 insertions(+), 18 deletions(-)
 create mode 100644 wren-ai-service/src/web/development.py

diff --git a/wren-ai-service/src/__main__.py b/wren-ai-service/src/__main__.py
index 0b3ba51cd..3d732e286 100644
--- a/wren-ai-service/src/__main__.py
+++ b/wren-ai-service/src/__main__.py
@@ -1,7 +1,5 @@
-import asyncio
 import logging
 import os
-import time
 from contextlib import asynccontextmanager
 
 import uvicorn
@@ -36,7 +34,6 @@ async def lifespan(app: FastAPI):
 
 app = FastAPI(lifespan=lifespan, redoc_url=None)
 
-app.include_router(routers.router, prefix="/v1")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -44,6 +41,11 @@ app.add_middleware(
     allow_methods=["*"],
     allow_headers=["*"],
 )
+app.include_router(routers.router, prefix="/v1")
+if os.getenv("ENV") == "dev":
+    from src.web import development
+
+    app.include_router(development.router, prefix="/dev")
 
 
 @app.exception_handler(Exception)
@@ -72,20 +74,6 @@ def health():
     return {"status": "ok"}
 
 
-@app.get("/dummy")
-async def dummy(sleep: int = 4, is_async: bool = True, should_sleep: bool = True):
-    """
-    Dummy endpoint to test async behavior by sleeping for several seconds
-    """
-    if should_sleep:
-        if is_async:
-            await asyncio.sleep(sleep)
-        else:
-            time.sleep(sleep)
-
-    return {"status": "dummy"}
-
-
 if __name__ == "__main__":
     server_host = os.getenv("WREN_AI_SERVICE_HOST") or "127.0.0.1"
     server_port = (
diff --git a/wren-ai-service/src/web/development.py b/wren-ai-service/src/web/development.py
new file mode 100644
index 000000000..78c3b9668
--- /dev/null
+++ b/wren-ai-service/src/web/development.py
@@ -0,0 +1,20 @@
+import asyncio
+import time
+
+from fastapi import APIRouter
+
+router = APIRouter()
+
+
+@router.get("/dummy")
+async def dummy(sleep: int = 4, is_async: bool = True, should_sleep: bool = True):
+    """
+    Dummy endpoint to test async behavior by sleeping for several seconds
+    """
+    if should_sleep:
+        if is_async:
+            await asyncio.sleep(sleep)
+        else:
+            time.sleep(sleep)
+
+    return {"status": "dummy"}
diff --git a/wren-ai-service/tests/locust/locustfile.py b/wren-ai-service/tests/locust/locustfile.py
index bd55dacef..a9264b3fb 100644
--- a/wren-ai-service/tests/locust/locustfile.py
+++ b/wren-ai-service/tests/locust/locustfile.py
@@ -230,4 +230,4 @@ class AskDetailsUser(FastHttpUser):
 class DummyUser(FastHttpUser):
     @task
     def dummy(self):
-        self.client.get(url="/dummy")
+        self.client.get(url="/dev/dummy")