|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +from __future__ import annotations |
| 8 | + |
| 9 | +import importlib |
| 10 | +import sys |
| 11 | +import tempfile |
| 12 | +import unittest |
| 13 | +from collections.abc import Callable |
| 14 | +from pathlib import Path |
| 15 | +from typing import get_origin |
| 16 | + |
| 17 | +from spdl.autoresearch._app._engine import _parse_engine_args |
| 18 | +from spdl.autoresearch._app._spec import ( |
| 19 | + _read_workflow_factory, |
| 20 | + _record_workflow_factory, |
| 21 | + _resolve_workflow, |
| 22 | +) |
| 23 | +from spdl.autoresearch._app._supervisor import ( |
| 24 | + _build_engine_command, |
| 25 | + _parse_supervisor_args, |
| 26 | +) |
| 27 | + |
| 28 | +__all__: list[str] = [] |
| 29 | + |
| 30 | + |
| 31 | +def _identity_factory(argv: list[str], workdir: Path | None) -> object: |
| 32 | + """A trivial factory used as a resolution target by the tests below.""" |
| 33 | + return (argv, workdir) |
| 34 | + |
| 35 | + |
| 36 | +class _ResolveWorkflowTest(unittest.TestCase): |
| 37 | + def test_module_factory_form(self) -> None: |
| 38 | + """_resolve_workflow imports module.path:factory_name and returns the callable.""" |
| 39 | + factory = _resolve_workflow(f"{__name__}:_identity_factory") |
| 40 | + self.assertIs(factory, _identity_factory) |
| 41 | + |
| 42 | + def test_empty_specifier_raises(self) -> None: |
| 43 | + """An empty string is rejected with ValueError, not silently importing.""" |
| 44 | + with self.assertRaises(ValueError): |
| 45 | + _resolve_workflow("") |
| 46 | + |
| 47 | + def test_malformed_specifier_raises(self) -> None: |
| 48 | + """A specifier with a colon but missing one half is rejected.""" |
| 49 | + for bad in (":factory", "module.path:", ":"): |
| 50 | + with self.subTest(bad=bad): |
| 51 | + with self.assertRaises(ValueError): |
| 52 | + _resolve_workflow(bad) |
| 53 | + |
| 54 | + def test_unknown_module_raises_import_error(self) -> None: |
| 55 | + """A non-existent module surfaces ModuleNotFoundError to the caller.""" |
| 56 | + with self.assertRaises(ModuleNotFoundError): |
| 57 | + _resolve_workflow("definitely.not.a.real.module:create") |
| 58 | + |
| 59 | + def test_missing_attribute_raises(self) -> None: |
| 60 | + """An existing module with a missing attribute raises AttributeError.""" |
| 61 | + with self.assertRaises(AttributeError): |
| 62 | + _resolve_workflow(f"{__name__}:does_not_exist") |
| 63 | + |
| 64 | + def test_short_name_lookup_misses_cleanly(self) -> None: |
| 65 | + """Short-name lookup raises LookupError when no entry point matches.""" |
| 66 | + with self.assertRaises(LookupError): |
| 67 | + _resolve_workflow("not_registered_workflow_xyz") |
| 68 | + |
| 69 | + |
| 70 | +class _WorkflowFactoryRecordTest(unittest.TestCase): |
| 71 | + def test_round_trip(self) -> None: |
| 72 | + """_record_workflow_factory followed by _read_workflow_factory returns the spec.""" |
| 73 | + with tempfile.TemporaryDirectory() as tmp: |
| 74 | + workdir = Path(tmp) |
| 75 | + _record_workflow_factory(workdir, "pkg.mod:factory") |
| 76 | + self.assertEqual(_read_workflow_factory(workdir), "pkg.mod:factory") |
| 77 | + |
| 78 | + def test_read_returns_none_when_missing(self) -> None: |
| 79 | + """_read_workflow_factory on a fresh workdir returns None instead of raising.""" |
| 80 | + with tempfile.TemporaryDirectory() as tmp: |
| 81 | + self.assertIsNone(_read_workflow_factory(Path(tmp))) |
| 82 | + |
| 83 | + def test_read_rejects_malformed_record(self) -> None: |
| 84 | + """Reading a malformed record file raises ValueError.""" |
| 85 | + with tempfile.TemporaryDirectory() as tmp: |
| 86 | + workdir = Path(tmp) |
| 87 | + (workdir / "workflow_factory.json").write_text("[]\n") |
| 88 | + with self.assertRaises(ValueError): |
| 89 | + _read_workflow_factory(workdir) |
| 90 | + |
| 91 | + def test_record_creates_workdir(self) -> None: |
| 92 | + """_record_workflow_factory creates the workdir if it does not yet exist.""" |
| 93 | + with tempfile.TemporaryDirectory() as tmp: |
| 94 | + workdir = Path(tmp) / "nested" |
| 95 | + _record_workflow_factory(workdir, "pkg.mod:factory") |
| 96 | + self.assertEqual(_read_workflow_factory(workdir), "pkg.mod:factory") |
| 97 | + |
| 98 | + |
| 99 | +class _ArgvSplitTest(unittest.TestCase): |
| 100 | + def test_supervisor_splits_at_double_dash(self) -> None: |
| 101 | + """Tokens after '--' are forwarded as the workflow tail, not parsed by the framework.""" |
| 102 | + ns, tail = _parse_supervisor_args( |
| 103 | + [ |
| 104 | + "/tmp/workdir", |
| 105 | + "--workflow", |
| 106 | + "pkg.mod:factory", |
| 107 | + "--max-concurrency", |
| 108 | + "5", |
| 109 | + "--", |
| 110 | + "--pipeline-script", |
| 111 | + "x.py", |
| 112 | + ] |
| 113 | + ) |
| 114 | + self.assertEqual(ns.workdir, "/tmp/workdir") |
| 115 | + self.assertEqual(ns.workflow, "pkg.mod:factory") |
| 116 | + self.assertEqual(ns.max_concurrency, 5) |
| 117 | + self.assertEqual(tail, ["--pipeline-script", "x.py"]) |
| 118 | + |
| 119 | + def test_supervisor_workdir_optional(self) -> None: |
| 120 | + """The supervisor accepts no workdir during initial config gathering.""" |
| 121 | + ns, tail = _parse_supervisor_args(["--workflow", "pkg.mod:factory"]) |
| 122 | + self.assertIsNone(ns.workdir) |
| 123 | + self.assertEqual(tail, []) |
| 124 | + |
| 125 | + def test_engine_requires_workflow_and_workdir(self) -> None: |
| 126 | + """The engine refuses to start without --workflow and --workdir.""" |
| 127 | + with self.assertRaises(SystemExit): |
| 128 | + _parse_engine_args(["--workdir", "/tmp/wd"]) |
| 129 | + with self.assertRaises(SystemExit): |
| 130 | + _parse_engine_args(["--workflow", "pkg.mod:factory"]) |
| 131 | + |
| 132 | + def test_engine_passes_tail_through(self) -> None: |
| 133 | + """The engine surfaces the workflow tail unchanged to the caller.""" |
| 134 | + ns, tail = _parse_engine_args( |
| 135 | + [ |
| 136 | + "--workflow", |
| 137 | + "pkg.mod:factory", |
| 138 | + "--workdir", |
| 139 | + "/tmp/wd", |
| 140 | + "--", |
| 141 | + "--build-command", |
| 142 | + "make", |
| 143 | + ] |
| 144 | + ) |
| 145 | + self.assertEqual(ns.workflow, "pkg.mod:factory") |
| 146 | + self.assertEqual(ns.workdir, "/tmp/wd") |
| 147 | + self.assertEqual(tail, ["--build-command", "make"]) |
| 148 | + |
| 149 | + def test_engine_max_concurrency_defaults_to_none(self) -> None: |
| 150 | + """Omitting --max-concurrency leaves ns.max_concurrency as None. |
| 151 | +
|
| 152 | + The engine treats this as "use the workflow-supplied default" so |
| 153 | + the WorkflowSpec.max_concurrency value is honored when the user |
| 154 | + does not override it on the CLI. |
| 155 | + """ |
| 156 | + ns, _ = _parse_engine_args( |
| 157 | + ["--workflow", "pkg.mod:factory", "--workdir", "/tmp/wd"] |
| 158 | + ) |
| 159 | + self.assertIsNone(ns.max_concurrency) |
| 160 | + |
| 161 | + def test_engine_max_concurrency_accepts_explicit_value(self) -> None: |
| 162 | + """An explicit --max-concurrency value is preserved as an int.""" |
| 163 | + ns, _ = _parse_engine_args( |
| 164 | + [ |
| 165 | + "--workflow", |
| 166 | + "pkg.mod:factory", |
| 167 | + "--workdir", |
| 168 | + "/tmp/wd", |
| 169 | + "--max-concurrency", |
| 170 | + "7", |
| 171 | + ] |
| 172 | + ) |
| 173 | + self.assertEqual(ns.max_concurrency, 7) |
| 174 | + |
| 175 | + |
| 176 | +class _EngineCommandTest(unittest.TestCase): |
| 177 | + def test_default_uses_spdl_autoresearch_engine(self) -> None: |
| 178 | + """Without an override, the engine prefix is 'spdl autoresearch engine'.""" |
| 179 | + cmd = _build_engine_command( |
| 180 | + engine_command_override=None, |
| 181 | + workflow_spec="pkg.mod:factory", |
| 182 | + workdir=Path("/tmp/wd"), |
| 183 | + framework_flags=["--max-concurrency", "3"], |
| 184 | + workflow_argv_tail=["--build-command", "make"], |
| 185 | + ) |
| 186 | + self.assertEqual( |
| 187 | + cmd, |
| 188 | + [ |
| 189 | + "spdl", |
| 190 | + "autoresearch", |
| 191 | + "engine", |
| 192 | + "--workflow", |
| 193 | + "pkg.mod:factory", |
| 194 | + "--workdir", |
| 195 | + "/tmp/wd", |
| 196 | + "--max-concurrency", |
| 197 | + "3", |
| 198 | + "--", |
| 199 | + "--build-command", |
| 200 | + "make", |
| 201 | + ], |
| 202 | + ) |
| 203 | + |
| 204 | + def test_override_replaces_prefix(self) -> None: |
| 205 | + """An --engine-command override replaces the default argv[0] prefix.""" |
| 206 | + cmd = _build_engine_command( |
| 207 | + engine_command_override="buck run //x:engine --", |
| 208 | + workflow_spec="pkg.mod:factory", |
| 209 | + workdir=Path("/tmp/wd"), |
| 210 | + framework_flags=[], |
| 211 | + workflow_argv_tail=[], |
| 212 | + ) |
| 213 | + self.assertEqual( |
| 214 | + cmd, |
| 215 | + [ |
| 216 | + "buck", |
| 217 | + "run", |
| 218 | + "//x:engine", |
| 219 | + "--", |
| 220 | + "--workflow", |
| 221 | + "pkg.mod:factory", |
| 222 | + "--workdir", |
| 223 | + "/tmp/wd", |
| 224 | + ], |
| 225 | + ) |
| 226 | + |
| 227 | + def test_no_tail_omits_double_dash(self) -> None: |
| 228 | + """An empty workflow tail does not append a stray '--'.""" |
| 229 | + cmd = _build_engine_command( |
| 230 | + engine_command_override=None, |
| 231 | + workflow_spec="pkg.mod:factory", |
| 232 | + workdir=Path("/tmp/wd"), |
| 233 | + framework_flags=[], |
| 234 | + workflow_argv_tail=[], |
| 235 | + ) |
| 236 | + self.assertNotIn("--", cmd[2:]) |
| 237 | + |
| 238 | + |
| 239 | +class _CoreWorkflowExportTest(unittest.TestCase): |
| 240 | + def test_workflow_spec_is_protocol(self) -> None: |
| 241 | + """``WorkflowSpec`` re-exported from core is a ``Protocol`` subclass. |
| 242 | +
|
| 243 | + ``Protocol`` subclasses are flagged with ``_is_protocol = True`` by |
| 244 | + the typing machinery; this guards against accidentally weakening |
| 245 | + ``WorkflowSpec`` to a regular class (which would silently change |
| 246 | + the runtime semantics for workflow authors). |
| 247 | + """ |
| 248 | + from spdl.autoresearch.core import WorkflowSpec |
| 249 | + |
| 250 | + self.assertTrue(getattr(WorkflowSpec, "_is_protocol", False)) |
| 251 | + |
| 252 | + def test_workflow_factory_is_callable_alias(self) -> None: |
| 253 | + """``WorkflowFactory`` re-exported from core is a ``Callable`` alias.""" |
| 254 | + from spdl.autoresearch.core import WorkflowFactory |
| 255 | + |
| 256 | + self.assertIs(get_origin(WorkflowFactory), Callable) |
| 257 | + |
| 258 | + |
| 259 | +class _MainImportTest(unittest.TestCase): |
| 260 | + def test_main_import_does_not_load_app(self) -> None: |
| 261 | + """Importing spdl.autoresearch.__main__ as a module is a no-op. |
| 262 | +
|
| 263 | + The framework dispatcher (under spdl.autoresearch._app) must |
| 264 | + NOT be transitively loaded by ``import |
| 265 | + spdl.autoresearch.__main__``. _app is reachable only when |
| 266 | + __main__.py runs as a script (i.e. via ``python -m |
| 267 | + spdl.autoresearch``), at which point ``__name__ == |
| 268 | + "__main__"`` and the lazy import inside the guard fires. |
| 269 | + """ |
| 270 | + removed = {} |
| 271 | + for mod_name in [ |
| 272 | + name |
| 273 | + for name in list(sys.modules) |
| 274 | + if name == "spdl.autoresearch.__main__" |
| 275 | + or name.startswith("spdl.autoresearch._app") |
| 276 | + ]: |
| 277 | + removed[mod_name] = sys.modules.pop(mod_name) |
| 278 | + self.addCleanup(sys.modules.update, removed) |
| 279 | + |
| 280 | + importlib.import_module("spdl.autoresearch.__main__") |
| 281 | + |
| 282 | + self.assertNotIn("spdl.autoresearch._app", sys.modules) |
| 283 | + self.assertNotIn("spdl.autoresearch._app._main", sys.modules) |
![meta-codesync[bot]](https://avatars.githubusercontent.com/in/1373969?v=4&size=40){kind=avatar}
0 commit comments