feat: adapting tests accordingly;

TogetherCrew · May 24, 2024 · b1b6506 · b1b6506
1 parent bd16ced
commit b1b6506
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/dags/hivemind_etl_helpers/tests/unit/test_mediawiki_extractor.py b/dags/hivemind_etl_helpers/tests/unit/test_mediawiki_extractor.py
@@ -31,7 +31,8 @@ def test_extract_from_valid_pages(self):
         test_pages = ["Python_(programming_language)", "OpenAI"]
         documents = self.extractor.extract(page_ids=test_pages)
         self.assertEqual(len(documents), len(mock_response))
-        self.mock_reader.load_data.assert_called_once_with(pages=test_pages)
+        self.mock_reader.load_data.assert_called_once_with(pages=test_pages,
+                                                           auto_suggest=False)
 
     def test_extract_no_pages(self):
         """
@@ -52,4 +53,5 @@ def test_handle_invalid_page_titles(self):
 
         documents = self.extractor.extract(page_ids=invalid_pages)
         self.assertEqual(len(documents), 0)
-        self.mock_reader.load_data.assert_called_with(pages=invalid_pages)
+        self.mock_reader.load_data.assert_called_with(pages=invalid_pages,
+                                                      auto_suggest=False)