update tests to set the payload in the metadata

nv-morpheus · Sep 23, 2024 · 4646eb2 · 4646eb2
1 parent 1d811c8
commit 4646eb2
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 0 deletions.
diff --git a/tests/llm/nodes/test_extractor_node.py b/tests/llm/nodes/test_extractor_node.py
@@ -39,6 +39,7 @@ def test_execute():
     df = cudf.DataFrame({"insects": insects.copy(), "mammals": mammals.copy(), "reptiles": reptiles.copy()})
     message = ControlMessage()
     message.payload(MessageMeta(df))
+    message.set_metadata("llm_message_meta", message.payload())
 
     task_dict = {"input_keys": ["mammals", "reptiles"]}
     node = ExtracterNode()

diff --git a/tests/llm/nodes/test_manual_extractor_node.py b/tests/llm/nodes/test_manual_extractor_node.py
@@ -48,6 +48,7 @@ def test_execute():
     df = cudf.DataFrame({"insects": insects.copy(), "mammals": mammals.copy(), "reptiles": reptiles.copy()})
     message = ControlMessage()
     message.payload(MessageMeta(df))
+    message.set_metadata("llm_message_meta", message.payload())
 
     task_dict = {"input_keys": ["insects"]}
     node = ManualExtracterNode(["mammals", "reptiles"])

diff --git a/tests/llm/task_handlers/test_simple_task_handler.py b/tests/llm/task_handlers/test_simple_task_handler.py
@@ -46,6 +46,7 @@ def test_try_handle(dataset_cudf: DatasetManager):
 
     message = ControlMessage()
     message.payload(MessageMeta(df))
+    message.set_metadata("llm_message_meta", message.payload())
 
     task_handler = SimpleTaskHandler(['reptiles'])