Add test for issue nv-morpheus#953

tests/test_multi_port_pipeline.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -16,27 +16,45 @@
 
 import pytest
 
+import cudf
+
 # When segment modules are imported, they're added to the module registry.
 # To avoid flake8 warnings about unused code, the noqa flag is used during import.
 import modules.multiplexer  # noqa: F401 # pylint: disable=unused-import
 from _utils.dataset_manager import DatasetManager
+from morpheus.config import Config
 from morpheus.pipeline.pipeline import Pipeline
 from morpheus.stages.general.multi_port_modules_stage import MultiPortModulesStage
 from morpheus.stages.input.in_memory_source_stage import InMemorySourceStage
 from morpheus.stages.output.in_memory_sink_stage import InMemorySinkStage
 
 
+def _run_pipeline(config: Config, source_df: cudf.DataFrame, module_conf: dict,
+                  stage_input_ports: list[str]) -> InMemorySinkStage:
+    pipe = Pipeline(config)
+
+    mux_stage = pipe.add_stage(
+        MultiPortModulesStage(config, module_conf, input_ports=stage_input_ports, output_ports=["output"]))
+
+    for x in range(len(stage_input_ports)):
+        source_stage = pipe.add_stage(InMemorySourceStage(config, [source_df.copy(deep=True)]))
+        pipe.add_edge(source_stage, mux_stage.input_ports[x])
+
+    sink_stage = pipe.add_stage(InMemorySinkStage(config))
+
+    pipe.add_edge(mux_stage, sink_stage)
+
+    pipe.run()
+
+    return sink_stage
+
+
 @pytest.mark.parametrize("source_count, expected_count", [(1, 1), (2, 2), (3, 3)])
-def test_multi_port_pipeline(config, dataset_cudf: DatasetManager, source_count, expected_count):
+def test_multi_port_pipeline(config: Config, dataset_cudf: DatasetManager, source_count, expected_count):
 
     filter_probs_df = dataset_cudf["filter_probs.csv"]
 
-    pipe = Pipeline(config)
-
-    input_ports = []
-    for x in range(source_count):
-        input_port = f"input_{x}"
-        input_ports.append(input_port)
+    input_ports = [f"input_{x}" for x in range(source_count)]
 
     multiplexer_module_conf = {
         "module_id": "multiplexer",
@@ -46,17 +64,29 @@ def test_multi_port_pipeline(config, dataset_cudf: DatasetManager, source_count,
         "output_port": "output"
     }
 
-    mux_stage = pipe.add_stage(
-        MultiPortModulesStage(config, multiplexer_module_conf, input_ports=input_ports, output_ports=["output"]))
+    sink_stage = _run_pipeline(config=config,
+                               source_df=filter_probs_df,
+                               module_conf=multiplexer_module_conf,
+                               stage_input_ports=input_ports)
+    assert len(sink_stage.get_messages()) == expected_count
 
-    for x in range(source_count):
-        source_stage = pipe.add_stage(InMemorySourceStage(config, [filter_probs_df]))
-        pipe.add_edge(source_stage, mux_stage.input_ports[x])
 
-    sink_stage = pipe.add_stage(InMemorySinkStage(config))
+def test_multi_port_pipeline_mis_config(config: Config, dataset_cudf: DatasetManager):
+    config_input_ports = ["input_0", "input_1"]
+    stage_input_ports = ["input_0", "input_1", "input_2"]
 
-    pipe.add_edge(mux_stage, sink_stage)
+    filter_probs_df = dataset_cudf["filter_probs.csv"]
 
-    pipe.run()
+    multiplexer_module_conf = {
+        "module_id": "multiplexer",
+        "namespace": "morpheus_test",
+        "module_name": "multiplexer",
+        "input_ports": config_input_ports,
+        "output_port": "output"
+    }
 
-    assert len(sink_stage.get_messages()) == expected_count
+    with pytest.raises(ValueError):
+        _run_pipeline(config=config,
+                      source_df=filter_probs_df,
+                      module_conf=multiplexer_module_conf,
+                      stage_input_ports=stage_input_ports)