Update post-processing tests

nv-morpheus · Dec 20, 2023 · 3821243 · 3821243
1 parent b85776c
commit 3821243
Showing 1 changed file with 7 additions and 14 deletions.
diff --git a/tests/examples/log_parsing/test_postprocessing.py b/tests/examples/log_parsing/test_postprocessing.py
@@ -26,9 +26,10 @@
 from morpheus.config import Config
 from morpheus.messages import InferenceMemory
 from morpheus.messages import MessageMeta
+from morpheus.messages import MultiInferenceMessage
 
 
-def build_post_proc_message(messages_mod, dataset_cudf: DatasetManager, log_test_data_dir: str):
+def build_post_proc_message(dataset_cudf: DatasetManager, log_test_data_dir: str):
     input_file = os.path.join(TEST_DIRS.validation_data_dir, 'log-parsing-validation-data-input.csv')
     input_df = dataset_cudf[input_file]
     meta = MessageMeta(input_df)
@@ -41,24 +42,16 @@ def build_post_proc_message(messages_mod, dataset_cudf: DatasetManager, log_test
         host_data = np.loadtxt(tensor_file, delimiter=',')
         tensors[tensor_name] = cp.asarray(host_data)
 
-    memory = InferenceMemory(count=5, **tensors)
-    return messages_mod.MultiPostprocLogParsingMessage(meta=meta,
-                                                       mess_offset=0,
-                                                       mess_count=count,
-                                                       memory=memory,
-                                                       offset=0,
-                                                       count=count)
+    memory = InferenceMemory(count=5, tensors=tensors)
+    return MultiInferenceMessage(meta=meta, mess_offset=0, mess_count=count, memory=memory, offset=0, count=count)
 
 
 @pytest.mark.use_python
-@pytest.mark.import_mod([
-    os.path.join(TEST_DIRS.examples_dir, 'log_parsing', 'messages.py'),
-    os.path.join(TEST_DIRS.examples_dir, 'log_parsing', 'postprocessing.py')
-])
+@pytest.mark.import_mod(os.path.join(TEST_DIRS.examples_dir, 'log_parsing', 'postprocessing.py'))
 def test_log_parsing_post_processing_stage(config: Config,
                                            dataset_cudf: DatasetManager,
                                            import_mod: typing.List[types.ModuleType]):
-    messages_mod, postprocessing_mod = import_mod
+    postprocessing_mod = import_mod
 
     model_vocab_file = os.path.join(TEST_DIRS.data_dir, 'bert-base-cased-vocab.txt')
     log_test_data_dir = os.path.join(TEST_DIRS.tests_data_dir, 'examples/log_parsing')
@@ -68,7 +61,7 @@ def test_log_parsing_post_processing_stage(config: Config,
                                                              vocab_path=model_vocab_file,
                                                              model_config_path=model_config_file)
 
-    post_proc_message = build_post_proc_message(messages_mod, dataset_cudf, log_test_data_dir)
+    post_proc_message = build_post_proc_message(dataset_cudf, log_test_data_dir)
     expected_df = dataset_cudf.pandas[os.path.join(log_test_data_dir, 'expected_out.csv')]
 
     out_meta = stage._postprocess(post_proc_message)