update: test cases with latest code updates!

TogetherCrew · Dec 20, 2023 · 4ff253e · 4ff253e
1 parent 838ce68
commit 4ff253e
Show file tree

Hide file tree

Showing 7 changed files with 51 additions and 26 deletions.
diff --git a/dags/hivemind_etl_helpers/tests/integration/test_discord_prepare_llama.py b/dags/hivemind_etl_helpers/tests/integration/test_discord_prepare_llama.py
@@ -57,7 +57,7 @@ def test_transform_two_data(self):
         data = {
             "type": 19,
             "author": "112",
-            "content": "mentioning <@113> <@114> <@101>",
+            "content": "mentioning <@113> <@114> <@&101>",
             "user_mentions": ["113", "114"],
             "role_mentions": ["101"],
             "reactions": [],

diff --git a/dags/hivemind_etl_helpers/tests/integration/test_discord_prepare_summary.py b/dags/hivemind_etl_helpers/tests/integration/test_discord_prepare_summary.py
@@ -27,7 +27,7 @@ def test_empty_data_prepare_without_per_date(self):
             thread_summary_docs,
             channel_summary_docs,
             day_summary_docs,
-        ) = discord_summary.prepare_summaries(guild_id)
+        ) = discord_summary.prepare_summaries(guild_id, summarization_prefix="")
 
         self.assertEqual(thread_summary_docs, [])
         self.assertEqual(channel_summary_docs, [])
@@ -47,7 +47,9 @@ def test_empty_data_prepare_with_from_date(self):
             thread_summary_docs,
             channel_summary_docs,
             day_summary_docs,
-        ) = discord_summary.prepare_summaries(guild_id, from_date=from_date)
+        ) = discord_summary.prepare_summaries(
+            guild_id, from_date=from_date, summarization_prefix=""
+        )
 
         self.assertEqual(thread_summary_docs, [])
         self.assertEqual(channel_summary_docs, [])
@@ -133,7 +135,9 @@ def test_some_data_prepare_with_from_date(self):
             thread_summary_docs,
             channel_summary_docs,
             day_summary_docs,
-        ) = discord_summary.prepare_summaries(guild_id, from_date=from_date)
+        ) = discord_summary.prepare_summaries(
+            guild_id, from_date=from_date, summarization_prefix=""
+        )
 
         # we had 2 days with 3 channels of each 1 thread
         self.assertEqual(len(thread_summary_docs), 6)
@@ -233,7 +237,9 @@ def test_some_data_prepare_after_from_date(self):
             thread_summary_docs,
             channel_summary_docs,
             day_summary_docs,
-        ) = discord_summary.prepare_summaries(guild_id, from_date=from_date)
+        ) = discord_summary.prepare_summaries(
+            guild_id, from_date=from_date, summarization_prefix=""
+        )
 
         self.assertEqual(thread_summary_docs, [])
         self.assertEqual(channel_summary_docs, [])

diff --git a/.../hivemind_etl_helpers/tests/integration/test_discord_summary_transform_daily_summaries.py b/.../hivemind_etl_helpers/tests/integration/test_discord_summary_transform_daily_summaries.py
@@ -23,9 +23,18 @@ def test_transform_daily_summary_to_document(self):
 
         # Expected output documents
         expected_documents = [
-            Document(text="Summary 1", metadata={"date": "2023-01-01"}),
-            Document(text="Summary 2", metadata={"date": "2023-01-02"}),
-            Document(text="Summary 3", metadata={"date": "2023-01-03"}),
+            Document(
+                text="Summary 1",
+                metadata={"date": "2023-01-01", "channel": None, "thread": None},
+            ),
+            Document(
+                text="Summary 2",
+                metadata={"date": "2023-01-02", "channel": None, "thread": None},
+            ),
+            Document(
+                text="Summary 3",
+                metadata={"date": "2023-01-03", "channel": None, "thread": None},
+            ),
         ]
 
         # Transform and get result

diff --git a/dags/hivemind_etl_helpers/tests/integration/test_gdrive_delete_records.py b/dags/hivemind_etl_helpers/tests/integration/test_gdrive_delete_records.py
@@ -48,7 +48,7 @@ def test_empty_inputs(self):
         no error should be raised
         """
         community_id = "1234"
-        table_name = "data_gdrive"
+        table_name = "gdrive"
 
         self.setUpDB(community_id=community_id)
         self.delete_previous_data(table_name)
@@ -66,7 +66,7 @@ def test_delete_one_record(self):
         no error should be raised
         """
         community_id = "1234"
-        table_name = "data_gdrive"
+        table_name = "gdrive"
 
         self.setUpDB(community_id=community_id)
         self.delete_previous_data(table_name)
@@ -88,9 +88,7 @@ def test_delete_one_record(self):
                 },
             )
         ]
-        pg_vector.save_documents(
-            documents,
-        )
+        pg_vector.save_documents(community_id, documents)
 
         delete_records(
             db_name=f"community_{community_id}",
@@ -105,7 +103,7 @@ def test_delete_multiple_record(self):
         no error should be raised
         """
         community_id = "1234"
-        table_name = "data_gdrive"
+        table_name = "gdrive"
 
         self.setUpDB(community_id=community_id)
         self.delete_previous_data(table_name)
@@ -137,9 +135,7 @@ def test_delete_multiple_record(self):
                 },
             ),
         ]
-        pg_vector.save_documents(
-            documents,
-        )
+        pg_vector.save_documents(community_id, documents)
 
         delete_records(
             db_name=f"community_{community_id}",
@@ -154,7 +150,7 @@ def test_delete_one_record_multiple_data_available(self):
         no error should be raised
         """
         community_id = "1234"
-        table_name = "data_gdrive"
+        table_name = "gdrive"
 
         self.setUpDB(community_id=community_id)
         self.delete_previous_data(table_name)
@@ -186,9 +182,7 @@ def test_delete_one_record_multiple_data_available(self):
                 },
             ),
         ]
-        pg_vector.save_documents(
-            documents,
-        )
+        pg_vector.save_documents(community_id, documents)
 
         # delete_records(
         #     db_name=f"community_{community_id}",

diff --git a/dags/hivemind_etl_helpers/tests/unit/test_discourse_summary_prepare_category.py b/dags/hivemind_etl_helpers/tests/unit/test_discourse_summary_prepare_category.py
@@ -17,9 +17,13 @@ def setUp(self):
     def test_prepare_daily_summaries_empty_data(self):
         self.setUp()
         forum_id = "12121221212"
+        forum_endpoint = "sample_endpoint"
 
         prepare_summaries = DiscourseSummary(
-            service_context=self.service_context, llm=self.mock_llm, forum_id=forum_id
+            service_context=self.service_context,
+            llm=self.mock_llm,
+            forum_id=forum_id,
+            forum_endpoint=forum_endpoint,
         )
 
         summaries, docs = prepare_summaries.prepare_daily_summaries(
@@ -33,9 +37,13 @@ def test_prepare_daily_summaries_empty_data(self):
     def test_prepare_daily_summaries_some_data(self):
         self.setUp()
         forum_id = "12121221212"
+        forum_endpoint = "sample_endpoint"
 
         prepare_summaries = DiscourseSummary(
-            service_context=self.service_context, llm=self.mock_llm, forum_id=forum_id
+            service_context=self.service_context,
+            llm=self.mock_llm,
+            forum_id=forum_id,
+            forum_endpoint=forum_endpoint,
         )
 
         category_summaries = {
@@ -69,9 +77,13 @@ def test_prepare_daily_summaries_some_data(self):
     def test_prepare_daily_summaries_some_data_check_documents(self):
         self.setUp()
         forum_id = "12121221212"
+        forum_endpoint = "sample_endpoint"
 
         prepare_summaries = DiscourseSummary(
-            service_context=self.service_context, llm=self.mock_llm, forum_id=forum_id
+            service_context=self.service_context,
+            llm=self.mock_llm,
+            forum_id=forum_id,
+            forum_endpoint=forum_endpoint,
         )
 
         category_summaries = {

diff --git a/dags/hivemind_etl_helpers/tests/unit/test_discourse_summary_prepare_topic.py b/dags/hivemind_etl_helpers/tests/unit/test_discourse_summary_prepare_topic.py
@@ -37,6 +37,7 @@ def test_prepare_topic_summaries_empty_data(self):
     def test_prepare_topic_summaries_some_data(self):
         self.setUp()
         forum_id = "12121221212"
+        forum_endpoint = "sample_endpoint"
 
         raw_data_grouped = []
         for i in range(2):
@@ -66,7 +67,10 @@ def test_prepare_topic_summaries_some_data(self):
             raw_data_grouped.append(data)
 
         prepare_summaries = DiscourseSummary(
-            service_context=self.service_context, llm=self.mock_llm, forum_id=forum_id
+            service_context=self.service_context,
+            llm=self.mock_llm,
+            forum_id=forum_id,
+            forum_endpoint=forum_endpoint,
         )
         summaries = prepare_summaries.prepare_topic_summaries(
             raw_data_grouped=raw_data_grouped,

diff --git a/dags/hivemind_etl_helpers/tests/unit/test_prepare_raw_message.py b/dags/hivemind_etl_helpers/tests/unit/test_prepare_raw_message.py
@@ -7,7 +7,7 @@
 
 class TestPrepareRawMessage(unittest.TestCase):
     def test_normal_raw_message_preparation(self):
-        msg_content = "Hello <@1234> and <@4321>"
+        msg_content = "Hello <@&1234> and <@4321>"
         roles_mention = {
             "1234": "Everyone",
         }