From 3a1e566b0fdfb699a32c5547c17de98fc2f5d09a Mon Sep 17 00:00:00 2001
From: Pao-Sheng Wang <paooap.oappao@gmail.com>
Date: Thu, 9 Jan 2025 11:50:45 +0800
Subject: [PATCH] fix: test cases for sql pair and deletion

---
 wren-ai-service/tests/data/pairs.json         | 14 +++++++++++
 .../pipelines/indexing/test_sql_pairs.py      | 25 ++++++++-----------
 .../indexing/test_sql_pairs_deletion.py       | 21 ++++++----------
 3 files changed, 31 insertions(+), 29 deletions(-)
 create mode 100644 wren-ai-service/tests/data/pairs.json

diff --git a/wren-ai-service/tests/data/pairs.json b/wren-ai-service/tests/data/pairs.json
new file mode 100644
index 000000000..5f5ab0967
--- /dev/null
+++ b/wren-ai-service/tests/data/pairs.json
@@ -0,0 +1,14 @@
+{
+    "test": [
+        {
+            "id": "1",
+            "question": "What is the book?",
+            "sql": "SELECT * FROM book"
+        },
+        {
+            "id": "2",
+            "question": "What is the author?",
+            "sql": "SELECT * FROM author"
+        }
+    ]
+}
\ No newline at end of file
diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py
index 23f819a73..bbce89413 100644
--- a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py
+++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py
@@ -2,7 +2,7 @@
 
 from src.config import settings
 from src.core.provider import DocumentStoreProvider
-from src.pipelines.indexing.sql_pairs import SqlPair, SqlPairs
+from src.pipelines.indexing import SqlPairs
 from src.providers import generate_components
 
 
@@ -17,12 +17,11 @@ async def test_sql_pairs_indexing_saving_to_document_store():
         recreate_index=True,
     )
 
-    sql_pairs = SqlPairs(**pipe_components["sql_pairs_indexing"])
+    sql_pairs = SqlPairs(
+        **pipe_components["sql_pairs_indexing"], sql_pairs_path="tests/data/pairs.json"
+    )
     await sql_pairs.run(
-        sql_pairs=[
-            SqlPair(sql="SELECT * FROM book", id="1", question="What is the book?"),
-            SqlPair(sql="SELECT * FROM author", id="2", question="What is the author?"),
-        ],
+        mdl_str='{"models": [{"properties": {"boilerplate": "test"}}]}',
         project_id="fake-id",
     )
 
@@ -46,20 +45,16 @@ async def test_sql_pairs_preparation_saving_to_document_store_with_multiple_proj
         recreate_index=True,
     )
 
-    sql_pairs = SqlPairs(**pipe_components["sql_pairs_indexing"])
+    sql_pairs = SqlPairs(
+        **pipe_components["sql_pairs_indexing"], sql_pairs_path="tests/data/pairs.json"
+    )
     await sql_pairs.run(
-        sql_pairs=[
-            SqlPair(sql="SELECT * FROM book", id="1", question="What is the book?"),
-            SqlPair(sql="SELECT * FROM author", id="2", question="What is the author?"),
-        ],
+        mdl_str='{"models": [{"properties": {"boilerplate": "test"}}]}',
         project_id="fake-id",
     )
 
     await sql_pairs.run(
-        sql_pairs=[
-            SqlPair(sql="SELECT * FROM book", id="1", question="What is the book?"),
-            SqlPair(sql="SELECT * FROM author", id="2", question="What is the author?"),
-        ],
+        mdl_str='{"models": [{"properties": {"boilerplate": "test"}}]}',
         project_id="fake-id-2",
     )
 
diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs_deletion.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs_deletion.py
index b654002aa..18226ab10 100644
--- a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs_deletion.py
+++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs_deletion.py
@@ -2,8 +2,7 @@
 
 from src.config import settings
 from src.core.provider import DocumentStoreProvider
-from src.pipelines.indexing.sql_pairs import SqlPair, SqlPairs
-from src.pipelines.indexing.sql_pairs_deletion import SqlPairsDeletion
+from src.pipelines.indexing import SqlPairs, SqlPairsDeletion
 from src.providers import generate_components
 
 
@@ -18,23 +17,17 @@ async def test_sql_pairs_deletion():
         recreate_index=True,
     )
 
-    sql_pairs = [
-        SqlPair(sql="SELECT * FROM book", id="1", question="What is the book?"),
-        SqlPair(sql="SELECT * FROM author", id="2", question="What is the author?"),
-    ]
-    sql_pairs = SqlPairs(**pipe_components["sql_pairs_indexing"])
+    sql_pairs = SqlPairs(
+        **pipe_components["sql_pairs_indexing"], sql_pairs_path="tests/data/pairs.json"
+    )
     await sql_pairs.run(
-        sql_pairs=sql_pairs,
+        mdl_str='{"models": [{"properties": {"boilerplate": "test"}}]}',
         project_id="fake-id",
     )
 
     sql_pairs_deletion = SqlPairsDeletion(**pipe_components["sql_pairs_deletion"])
-    await sql_pairs_deletion.run(
-        id="fake-id-2", sql_pair_ids=[sql_pair.id for sql_pair in sql_pairs]
-    )
+    await sql_pairs_deletion.run(id="fake-id-2", sql_pair_ids=["1", "2"])
     assert await store.count_documents() == 2
 
-    await sql_pairs_deletion.run(
-        id="fake-id", sql_pair_ids=[sql_pair.id for sql_pair in sql_pairs]
-    )
+    await sql_pairs_deletion.run(id="fake-id", sql_pair_ids=["1", "2"])
     assert await store.count_documents() == 0