existing e2e test for graphrag sorted

Author: serengil

tests/e2e/test_graphrag_e2e.py

@@ -60,7 +60,7 @@ def test_graphrag_happy_path(
     )
 
     llm.invoke.assert_called_once_with(
-        """Context:
+        input="""Context:
 <Record node={'question': 'In 1953 Watson & Crick built a model of the molecular structure of this, the gene-carrying substance'}>
 <Record node={'question': 'This organ removes excess glucose from the blood & stores it as glycogen'}>
 
@@ -72,7 +72,7 @@ def test_graphrag_happy_path(
 
 Answer:
 """,
-        None,
+        message_history=None,
         system_instruction="Answer the user question using the provided context.",
     )
     assert isinstance(result, RagResultModel)
@@ -152,8 +152,8 @@ def test_graphrag_happy_path_with_neo4j_message_history(
                 system_instruction=first_invocation_system_instruction,
             ),
             call(
-                second_invocation,
-                message_history.messages,
+                input=second_invocation,
+                message_history=message_history.messages,
                 system_instruction="Answer the user question using the provided context.",
             ),
         ]
@@ -190,7 +190,7 @@ def test_graphrag_happy_path_return_context(
     )
 
     llm.invoke.assert_called_once_with(
-        """Context:
+        input="""Context:
 <Record node={'question': 'In 1953 Watson & Crick built a model of the molecular structure of this, the gene-carrying substance'}>
 <Record node={'question': 'This organ removes excess glucose from the blood & stores it as glycogen'}>
 
@@ -202,7 +202,7 @@ def test_graphrag_happy_path_return_context(
 
 Answer:
 """,
-        None,
+        message_history=None,
         system_instruction="Answer the user question using the provided context.",
     )
     assert isinstance(result, RagResultModel)
@@ -236,7 +236,7 @@ def test_graphrag_happy_path_examples(
     )
 
     llm.invoke.assert_called_once_with(
-        """Context:
+        input="""Context:
 <Record node={'question': 'In 1953 Watson & Crick built a model of the molecular structure of this, the gene-carrying substance'}>
 <Record node={'question': 'This organ removes excess glucose from the blood & stores it as glycogen'}>
 
@@ -248,7 +248,7 @@ def test_graphrag_happy_path_examples(
 
 Answer:
 """,
-        None,
+        message_history=None,
         system_instruction="Answer the user question using the provided context.",
     )
     assert result.answer == "some text"