Test fixed now

Knox-AAU · Dec 7, 2023 · 3b4b89d · 3b4b89d
1 parent 7c32e65
commit 3b4b89d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git 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
@@ -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: <Barack Obama, married, Michelle Obama> and <Michelle Obama, married, Barack Obama>"
+                        "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: <Peter, sibling, Marianne> and <Marianne, sibling, Peter>"
+                        "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: <Barack Obama, married, Michelle Obama> and <Michelle Obama, married, Barack Obama>"
+                            "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">'
                         }
                     ],
                 },