From cd9081dad12155a96b4a0c2a630b3fdeebff3ba4 Mon Sep 17 00:00:00 2001 From: Michael Demoret Date: Fri, 25 Oct 2024 14:19:47 -0400 Subject: [PATCH] Fixing up the router tests --- python/morpheus/morpheus/_lib/stages/module.cpp | 9 ++++----- tests/morpheus/stages/test_router_stage_pipe.py | 9 +++++---- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/python/morpheus/morpheus/_lib/stages/module.cpp b/python/morpheus/morpheus/_lib/stages/module.cpp index 35ddc19d8a..6fb855e0c5 100644 --- a/python/morpheus/morpheus/_lib/stages/module.cpp +++ b/python/morpheus/morpheus/_lib/stages/module.cpp @@ -47,11 +47,10 @@ #include // for trace_activity, decay_t #include // for path -#include -#include // for shared_ptr, allocator -#include // for operator<<, basic_ostringstream -#include // for string -#include // for vector +#include // for shared_ptr, allocator +#include // for operator<<, basic_ostringstream +#include // for string +#include // for vector namespace morpheus { namespace py = pybind11; diff --git a/tests/morpheus/stages/test_router_stage_pipe.py b/tests/morpheus/stages/test_router_stage_pipe.py index 8d840f203e..861bb44b6d 100644 --- a/tests/morpheus/stages/test_router_stage_pipe.py +++ b/tests/morpheus/stages/test_router_stage_pipe.py @@ -29,8 +29,8 @@ from morpheus.stages.preprocess.deserialize_stage import DeserializeStage -@pytest.mark.parametrize("is_runnable", [True, False]) -def test_router_stage_pipe(config, filter_probs_df, is_runnable: bool): +@pytest.mark.parametrize("processing_engines", [0, 4]) +def test_router_stage_pipe(config, filter_probs_df, processing_engines: bool): keys = ["odd", "even"] @@ -44,7 +44,8 @@ def determine_route_fn(_: ControlMessage): pipe = Pipeline(config) source = pipe.add_stage(InMemorySourceStage(config, dataframes=[filter_probs_df], repeat=5)) deserialize = pipe.add_stage(DeserializeStage(config)) - router_stage = pipe.add_stage(RouterStage(config, keys=keys, key_fn=determine_route_fn, is_runnable=is_runnable)) + router_stage = pipe.add_stage( + RouterStage(config, keys=keys, key_fn=determine_route_fn, processing_engines=processing_engines)) sink1 = pipe.add_stage(InMemorySinkStage(config)) sink2 = pipe.add_stage(InMemorySinkStage(config)) @@ -101,7 +102,7 @@ def determine_route_fn(_: ControlMessage): pipe = Pipeline(config) source = pipe.add_stage(InMemoryDataGenStage(config, data_source=source_fn, output_data_type=ControlMessage)) - router_stage = pipe.add_stage(RouterStage(config, keys=keys, key_fn=determine_route_fn, is_runnable=True)) + router_stage = pipe.add_stage(RouterStage(config, keys=keys, key_fn=determine_route_fn, processing_engines=10)) sink1 = pipe.add_stage(InMemorySinkStage(config)) sink2 = pipe.add_stage(InMemorySinkStage(config))