diff --git a/test/test_relation_extraction/test_llama_relation_extractor.py b/test/test_relation_extraction/test_llama_relation_extractor.py index 1b531ea..021e004 100644 --- a/test/test_relation_extraction/test_llama_relation_extractor.py +++ b/test/test_relation_extraction/test_llama_relation_extractor.py @@ -29,7 +29,7 @@ def test_handle_post_request_raises_exception_if_parse_fail(self, mock_extract_s @mock.patch('relation_extraction.multilingual.main.parse_data') @mock.patch('relation_extraction.ontology_messenger.OntologyMessenger.send_request') def test_handle_post_request_raises_exception_if_prompt_llm_fail(self, mock_extract_specific_relations, mock_parse_data, mock_prompt_llm): - mock_extract_specific_relations.return_value = [] + mock_extract_specific_relations.return_value = ["relation1"] mock_parse_data.return_value = [] mock_prompt_llm.side_effect = Exception() @@ -39,16 +39,16 @@ def test_handle_post_request_raises_exception_if_prompt_llm_fail(self, mock_extr mock_extract_specific_relations.assert_called_once() mock_parse_data.assert_called_once() - mock_prompt_llm.assert_called_once() + mock_prompt_llm.assert_called() @mock.patch('relation_extraction.knowledge_graph_messenger.KnowledgeGraphMessenger.send_request') @mock.patch('relation_extraction.multilingual.llm_messenger.LLMMessenger.prompt_llm') @mock.patch('relation_extraction.multilingual.main.parse_data') @mock.patch('relation_extraction.ontology_messenger.OntologyMessenger.send_request') def test_handle_post_request_raises_exception_if_send_to_db_fail(self, mock_extract_specific_relations, mock_parse_data, mock_prompt_llm, mock_send_to_db): - mock_extract_specific_relations.return_value = [] + mock_extract_specific_relations.return_value = ["relation1"] mock_parse_data.return_value = [] - mock_prompt_llm.return_value = {} + mock_prompt_llm.return_value = [] mock_send_to_db.side_effect = Exception() data = dict() @@ -57,7 +57,7 @@ def test_handle_post_request_raises_exception_if_send_to_db_fail(self, mock_extr mock_extract_specific_relations.assert_called_once() mock_parse_data.assert_called_once() - mock_prompt_llm.assert_called_once() + mock_prompt_llm.assert_called() mock_send_to_db.assert_called_once() class TestParseData(unittest.TestCase): diff --git a/test/test_relation_extraction/test_llm_messenger.py b/test/test_relation_extraction/test_llm_messenger.py index 4301025..cdec19b 100644 --- a/test/test_relation_extraction/test_llm_messenger.py +++ b/test/test_relation_extraction/test_llm_messenger.py @@ -16,7 +16,7 @@ def test_process_message(self): { "choices": [ { - "text":"[INST] Barack Obama is married to Michelle Obama. [/INST] In this sentence the triples are: and " + "text":'[INST] Barack Obama is married to Michelle Obama. [/INST] In this sentence the triples are: <"Barack Obama", married, "Michelle Obama"> and <"Michelle Obama", married, "Barack Obama">' } ], "expected": [ @@ -35,7 +35,7 @@ def test_process_message(self): { "choices": [ { - "text":"[INST] Peter and Marianne has the same mother. [/INST] In this sentence the triples are: and " + "text":'[INST] Peter and Marianne has the same mother. [/INST] In this sentence the triples are: <"Peter", sibling, "Marianne"> and <"Marianne", sibling, "Peter">' } ], "expected":[ @@ -136,7 +136,7 @@ def test_prompt_llm(self, mock_check_validity, mock_process_message, mock_send_r "response": { "choices": [ { - "text":"[INST] Barack Obama is married to Michelle Obama. [/INST] In this sentence the triples are: and " + "text":'[INST] Barack Obama is married to Michelle Obama. [/INST] In this sentence the triples are: <"Barack Obama", married, "Michelle Obama"> and <"Michelle Obama", married, "Barack Obama">' } ], },