diff --git a/dags/hivemind_etl_helpers/tests/integration/test_telegram_chats.py b/dags/hivemind_etl_helpers/tests/integration/test_telegram_chats.py index 8a878c2b..69a983c1 100644 --- 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"), + ], + )