Robert Wojciechowski
3058b11ff2
- Remove DistributedListSplitter and DistributedListCollector nodes - Remove /distributed/init_list_queue and /distributed/request_list_item endpoints - Clean up all references from prompt_transform, queue_orchestration, constants, and tests - Refactor and reorganize distributed modules for quality improvements
72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
import importlib.util
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
def _load_distributed_module():
|
|
module_path = Path(__file__).resolve().parents[3] / "distributed.py"
|
|
package_name = "dist_distributed_entry_testpkg"
|
|
calls = {
|
|
"build_node_mappings": 0,
|
|
"initialize_runtime": 0,
|
|
"initialize_runtime_args": [],
|
|
}
|
|
|
|
for mod_name in list(sys.modules):
|
|
if mod_name == package_name or mod_name.startswith(f"{package_name}."):
|
|
del sys.modules[mod_name]
|
|
|
|
root_pkg = types.ModuleType(package_name)
|
|
root_pkg.__path__ = []
|
|
sys.modules[package_name] = root_pkg
|
|
|
|
bootstrap_pkg = types.ModuleType(f"{package_name}.bootstrap")
|
|
bootstrap_pkg.__path__ = []
|
|
sys.modules[f"{package_name}.bootstrap"] = bootstrap_pkg
|
|
|
|
entrypoint_module = types.ModuleType(f"{package_name}.bootstrap.entrypoint")
|
|
|
|
def _build_node_mappings():
|
|
calls["build_node_mappings"] += 1
|
|
return ({"MockNode": object}, {"MockNode": "Mock Node"})
|
|
|
|
def _initialize_runtime(prompt_server=None):
|
|
calls["initialize_runtime"] += 1
|
|
calls["initialize_runtime_args"].append(prompt_server)
|
|
|
|
entrypoint_module.build_node_mappings = _build_node_mappings
|
|
entrypoint_module.initialize_runtime = _initialize_runtime
|
|
sys.modules[f"{package_name}.bootstrap.entrypoint"] = entrypoint_module
|
|
|
|
spec = importlib.util.spec_from_file_location(f"{package_name}.distributed", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec is not None and spec.loader is not None
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module, calls
|
|
|
|
|
|
class DistributedEntryTests(unittest.TestCase):
|
|
def test_import_builds_mappings_without_runtime_bootstrap(self):
|
|
module, calls = _load_distributed_module()
|
|
|
|
self.assertIn("MockNode", module.NODE_CLASS_MAPPINGS)
|
|
self.assertEqual(module.NODE_DISPLAY_NAME_MAPPINGS["MockNode"], "Mock Node")
|
|
self.assertEqual(calls["build_node_mappings"], 1)
|
|
self.assertEqual(calls["initialize_runtime"], 0)
|
|
|
|
def test_initialize_runtime_delegates_to_bootstrap_entrypoint(self):
|
|
module, calls = _load_distributed_module()
|
|
sentinel = object()
|
|
|
|
module.initialize_runtime(sentinel)
|
|
|
|
self.assertEqual(calls["initialize_runtime"], 1)
|
|
self.assertEqual(calls["initialize_runtime_args"], [sentinel])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|