Updated test case assertions

Querent-ai · Apr 13, 2024 · ad1381b · ad1381b
1 parent 21e1d4d
commit ad1381b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/tests/ingestors/test_html_ingestor.py b/tests/ingestors/test_html_ingestor.py
@@ -39,7 +39,7 @@ async def poll_and_print():
             if ingested != "" or ingested is not None:
                 counter += 1
         # 1 extra IngestedToken signifying end of file
-        assert counter == 17
+        assert counter == 18
 
     await poll_and_print()
 

diff --git a/tests/llm_tests/mock_llm_test.py b/tests/llm_tests/mock_llm_test.py
@@ -34,7 +34,7 @@ async def process_tokens(self, data: IngestedTokens):
         # can set the state of the LLM using the set_state method
         # The state of the LLM is stored in the state attribute of the LLM
         # The state of the LLM is published to subscribers of the LLM
-        current_state = EventState(EventType.Graph, 1.0, "anything", "dummy.txt")
+        current_state = EventState(EventType.Graph, 1.0, "anything", "dummy.txt", doc_source="file://folder_path")
         await self.set_state(new_state=current_state)
 
     async def process_code(self, data: IngestedCode):
@@ -101,5 +101,5 @@ def handle_event(self, event_type: EventType, event_state: EventState):
 
 async def terminate_querent(llm_mocker: MockLLMEngine):
     await asyncio.sleep(60)
-    event_state = EventState(event_type=EventType.Graph, payload = "terminate", timestamp=1.0, file="dummy.txt")
+    event_state = EventState(event_type=EventType.Graph, payload = "terminate", timestamp=1.0, file="dummy.txt", doc_source="file://folder_path")
     await llm_mocker.set_state(event_state)