even more tweaks

CrayLabs · Oct 29, 2024 · 29fe4c3 · 29fe4c3
1 parent 616d816
commit 29fe4c3
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/tests/dragon_wlm/test_request_dispatcher.py b/tests/dragon_wlm/test_request_dispatcher.py
@@ -53,6 +53,9 @@
 from smartsim._core.mli.comm.channel.dragon_channel import DragonCommChannel
 from smartsim._core.mli.comm.channel.dragon_fli import DragonFLIChannel
 from smartsim._core.mli.comm.channel.dragon_util import create_local
+from smartsim._core.mli.infrastructure.control.dragon_util import (
+    function_as_dragon_proc,
+)
 from smartsim._core.mli.infrastructure.control.request_dispatcher import (
     RequestBatch,
     RequestDispatcher,
@@ -72,9 +75,8 @@
 from smartsim._core.mli.message_handler import MessageHandler
 from smartsim.log import get_logger
 
-from . import conftest
-from .utils import msg_pump
-from .utils.channel import FileSystemCommChannel
+from tests.dragon_wlm.utils.msg_pump import mock_messages
+from tests.dragon_wlm.utils.channel import FileSystemCommChannel
 
 logger = get_logger(__name__)
 
@@ -148,8 +150,8 @@ def test_request_dispatcher(
         callback_channel = DragonCommChannel.from_local()
         channels.append(callback_channel)
 
-        process = conftest.function_as_dragon_proc(
-            msg_pump.mock_messages,
+        process = function_as_dragon_proc(
+            mock_messages,
             [
                 worker_queue.descriptor,
                 backbone_fs.descriptor,