fix: test case updated with latest code structure!

TogetherCrew · Oct 21, 2024 · c5f6227 · c5f6227
1 parent c563c19
commit c5f6227
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/dags/hivemind_etl_helpers/tests/integration/test_telegram_chats.py b/dags/hivemind_etl_helpers/tests/integration/test_telegram_chats.py
@@ -51,8 +51,8 @@ def test_extract_chats_single_chat(self):
                 )
             )
 
-        chat_ids = self.tc_chats.extract_chats()
-        self.assertEqual(chat_ids, ["100000"])
+        chat_infos = self.tc_chats.extract_chats()
+        self.assertEqual(chat_infos, [("100000", "test chat")])
 
     def test_extract_chats_multiple_chats(self):
         neo4j_driver = self.tc_chats._connection.neo4j_driver
@@ -106,4 +106,11 @@ def test_extract_chats_multiple_chats(self):
             )
 
         chat_ids = self.tc_chats.extract_chats()
-        self.assertEqual(chat_ids, ["100001", "100002", "100003"])
+        self.assertEqual(
+            chat_ids,
+            [
+                ("100001", "test chat"),
+                ("100002", "test chat 2"),
+                ("100003", "test chat 3"),
+            ],
+        )