diff --git a/tests/examples/log_parsing/test_postprocessing.py b/tests/examples/log_parsing/test_postprocessing.py index c1d4b2c701..0446d2d2e2 100644 --- 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)