Merge pull request #28 from FacerAin/feat/test

Add test code and coverage

Makefile

@@ -4,6 +4,7 @@ style: set-style-dep set-style
 setup: set-precommit set-style-dep set-test-dep set-git set-dev
 test: set-test-dep set-test
 run: set-dev run-uvicorn
+coverage: set-coverage
 
 
 ##### basic #####
@@ -24,7 +25,7 @@ set-dev:
 	pip3 install -r requirements.txt
 
 set-test:
-	python3 -m pytest tests/
+	python3 -m pytest tests/ -m "not live"
 
 set-style:
 	black --config pyproject.toml .
@@ -51,3 +52,6 @@ clean-test:
 
 run-uvicorn:
 	uvicorn app.main:app
+
+set-coverage:
+	python3 -m pytest --cov-report term-missing --cov=app tests/ -m "not live"

app/agent/agent.py

@@ -10,6 +10,7 @@
 from app.agent.parser import CustomAgentOutputParser
 from app.agent.prompts import retriever_prompt_template, system_message
 from app.agent.retriever import PineconeRetriever
+from app.agent.tools import get_meal_info, get_today_date
 from app.core.config import settings
 
 set_llm_cache(InMemoryCache())
@@ -30,7 +31,7 @@ def __init__(self):
             ),
             Tool(
                 name="cafeterial_menu",
-                func=simple_meal_info,
+                func=get_meal_info,
                 description="If a user is looking for campus cafeterial menu information, use this information.",
             ),
             Tool(
@@ -53,19 +54,6 @@ def run(self, query):
         return response
 
 
-def simple_meal_info(query):
-    return """
-If a user is looking for campus cafeterial menu information, use the link below. You should directly check the meal information from the link below.
-경희대학교 학생 식당: https://www.khu.ac.kr/kor/forum/list.do?type=RESTAURANT&category=INTL&page=1
-경희대학교 제 2기숙사 식당: https://dorm2.khu.ac.kr/40/4050.kmc
-"""
-
-
-def get_today_date(query):
-    now = datetime.datetime.now()
-    return str(now.strftime("%Y-%m-%d"))
-
-
 class RetrieverAgent:
     def __init__(self, index_name: str = "khugpt") -> None:
         self.llm = ChatOpenAI(model_name="gpt-4-1106-preview", temperature=0, openai_api_key=settings.OPENAI_API_KEY)

app/agent/tools.py

@@ -0,0 +1,14 @@
+import datetime
+
+
+def get_meal_info(query: str = ""):
+    return """
+If a user is looking for campus cafeterial menu information, use the link below. You should directly check the meal information from the link below.
+경희대학교 학생 식당: https://www.khu.ac.kr/kor/forum/list.do?type=RESTAURANT&category=INTL&page=1
+경희대학교 제 2기숙사 식당: https://dorm2.khu.ac.kr/40/4050.kmc
+"""
+
+
+def get_today_date(query: str = ""):
+    now = datetime.datetime.now()
+    return str(now.strftime("%Y-%m-%d"))

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+markers =
+    live: For testing actual API interactions (deselect with '-m "not live"')

tests/test_client.py

@@ -0,0 +1,44 @@
+import pytest
+from fastapi.testclient import TestClient
+
+from app.agent import ExecutorAgent
+from app.agent.retriever import PineconeRetriever
+from app.main import app
+
+client = TestClient(app)
+
+
+def test_completion(monkeypatch):
+    def mockreturn(a, b):
+        return "Hello"
+
+    monkeypatch.setattr(ExecutorAgent, "run", mockreturn)
+    req_body = {"query": "Hi"}
+    response = client.post("/api/v1/chat/completion", json=req_body)
+    assert response.status_code == 200
+    assert "answer" in response.json()
+    assert "Hello" == response.json()["answer"]
+
+
+def test_similarity_search(monkeypatch):
+    def mockreturn(a, b):
+        return "Hello"
+
+    monkeypatch.setattr(PineconeRetriever, "get_relevant_doc_string", mockreturn)
+    response = client.get("/api/v1/chat/similarity-search?query=Hi")
+    assert response.status_code == 200
+    assert isinstance(response.text, str)
+    assert '"Hello"' == response.text
+
+
+@pytest.mark.live
+def test_live_completion():
+    req_body = {"query": "Hi"}
+    response = client.post("/api/v1/chat/completion", json=req_body)
+    assert response.status_code == 200
+
+
+@pytest.mark.live
+def test_live_similarity_search():
+    response = client.get("/api/v1/chat/similarity-search?query=Hi")
+    assert response.status_code == 200

tests/test_tools.py

@@ -0,0 +1,11 @@
+import datetime
+
+from app.agent.tools import get_meal_info, get_today_date
+
+
+def test_get_today_date():
+    assert get_today_date() == str(datetime.datetime.now().strftime("%Y-%m-%d"))
+
+
+def test_get_meal_info():
+    assert "https://dorm2.khu.ac.kr/40/4050.kmc" in get_meal_info()