DH-5080/update the tests

Dataherald · Dec 15, 2023 · 035d875 · 035d875
1 parent 019c696
commit 035d875
Showing 1 changed file with 6 additions and 8 deletions.
diff --git a/dataherald/tests/sql_generator/test_generator.py b/dataherald/tests/sql_generator/test_generator.py
@@ -5,7 +5,7 @@
 from dataherald.config import System
 from dataherald.sql_database.models.types import DatabaseConnection
 from dataherald.sql_generator import SQLGenerator
-from dataherald.types import Question, Response
+from dataherald.types import Prompt, SQLGeneration
 
 
 class TestGenerator(SQLGenerator):
@@ -15,14 +15,12 @@ def __init__(self, system: System):
     @override
     def generate_response(
         self,
-        user_question: Question,
+        user_question: Prompt,
         database_connection: DatabaseConnection,
         context: List[dict] = None,  # noqa: ARG002
-        generate_csv: bool = None,
-    ) -> Response:
-        return Response(
+    ) -> SQLGeneration:
+        return SQLGeneration(
             question_id="651f2d76275132d5b65175eb",
-            response="Foo response",
-            sql_query="bar",
-            generate_csv=None,
+            sql="Foo response",
+            status="bar",
         )