"""Offline tests of probe_orchestrator.py: a virtual clock and a fake Azure. Nothing is sent anywhere. python -m pytest -q test_orchestrator.py """ import importlib.util import os import types import zipfile import pytest HERE = os.path.dirname(os.path.abspath(__file__)) class FakeAzure: """ARM, Blob and Log Analytics as the orchestrator calls them. Each call takes one virtual second.""" def __init__(self, t, *, run_ends_after=None, start_ok=True, dealloc_works=True, state_reads_fail=False): self.t = t self.run_ends_after = run_ends_after # virtual seconds after submission; None = never ends self.start_ok, self.dealloc_works, self.state_reads_fail = start_ok, dealloc_works, state_reads_fail self.power = "VM deallocated" self.blobs, self.calls = {}, [] self.run_submitted = None self.RG, self.VM, self.ACCOUNT, self.CONTAINER, self.ARM = "rg", "vm", "acct", "c", "https://management.azure.com" import urllib.parse self.urllib = types.SimpleNamespace(parse=urllib.parse) def _tick(self): self.t[0] += 1 def vmpath(self): return "/vm" def token(self, res): return "tok" def blob_url(self, name): return f"blob://{name}" def arm(self, method, path, body=None, api=None): self._tick() self.calls.append((method, path)) if path.endswith("/instanceView"): if self.state_reads_fail: raise ConnectionError("no answer") return 200, {"statuses": [{"displayStatus": "Provisioning succeeded"}, {"displayStatus": self.power}]}, {} if path.endswith("/start"): if not self.start_ok: return 409, {"error": "conflict"}, {} self.power = "VM running" return 202, {}, {"Azure-AsyncOperation": "op://start"} if path.endswith("/deallocate"): if self.dealloc_works: self.power = "VM deallocated" return 202, {}, {"Azure-AsyncOperation": "op://dealloc"} if path.endswith("/runCommand"): self.run_submitted = self.t[0] return 202, {}, {"Azure-AsyncOperation": "op://run"} raise AssertionError(path) def http(self, method, url, data=None, headers=None, raw=False, timeout=120): self._tick() self.calls.append((method, url)) if url == "op://start": return 200, {"status": "Succeeded"}, {} if url == "op://dealloc": return 200, {"status": "Succeeded" if self.dealloc_works else "InProgress"}, {} if url == "op://run": if self.run_ends_after is not None and self.t[0] - self.run_submitted >= self.run_ends_after: return 200, {"status": "Succeeded", "properties": {"output": {"value": [{"message": "session done"}]}}}, {} return 200, {"status": "InProgress"}, {} if url.startswith("blob://"): name = url[len("blob://"):] if method == "PUT": self.blobs[name] = data return 201, b"", {} return (200, self.blobs[name], {}) if name in self.blobs else (404, b"", {}) if "loganalytics" in url: return 200, {"tables": [{"rows": [["t", "vm", 1200, "start"], ["t", "vm", 1201, "end"], ["t", "vm", 1100, "upload"]]}]}, {} raise AssertionError(url) @pytest.fixture def orch(tmp_path): spec = importlib.util.spec_from_file_location("probe_orchestrator_under_test", os.path.join(HERE, "probe_orchestrator.py")) mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) t = [0.0] mod.clock = lambda: t[0] mod.sleep = lambda s: t.__setitem__(0, t[0] + max(1.0, s)) rel = tmp_path / "release" rel.mkdir() with zipfile.ZipFile(rel / "marianos-scraper-2.2.16.zip", "w") as z: z.writestr("scraper/infra/run_probe.ps1", "launcher") args = ["--version", "2.2.16", "--prior-runs", "2", "--max-requests", "11", "--runs-dir", str(tmp_path / "runs"), "--release-dir", str(rel), "--session-deadline", "3600"] return mod, t, args def test_completed_session_exits_0_after_confirmed_deallocation(orch): mod, t, args = orch mod.az = FakeAzure(t, run_ends_after=600) assert mod.main(args) == 0 assert mod.az.power == "VM deallocated" and ("POST", "/vm/deallocate") in mod.az.calls def test_a_session_that_never_ends_is_cut_at_the_deadline_and_the_vm_is_deallocated(orch): mod, t, args = orch mod.az = FakeAzure(t, run_ends_after=None) # the VM never reports the end code = mod.main(args) assert code == 2 and mod.az.power == "VM deallocated" dealloc_at = next(i for i, c in enumerate(mod.az.calls) if c == ("POST", "/vm/deallocate")) assert dealloc_at > 0 and t[0] < 3600 + 60 # stopped waiting at the deadline, then deallocated def test_a_failed_start_still_deallocates_and_confirms(orch): mod, t, args = orch mod.az = FakeAzure(t, start_ok=False) assert mod.main(args) == 2 and ("POST", "/vm/deallocate") in mod.az.calls def test_deallocation_that_cannot_be_confirmed_is_retried_bounded_and_exits_3(orch, capsys): mod, t, args = orch mod.az = FakeAzure(t, run_ends_after=600, dealloc_works=False) code = mod.main(args) out = capsys.readouterr().out assert code == 3 and "VM NOT CONFIRMED DEALLOCATED" in out and "az vm deallocate" in out assert sum(1 for c in mod.az.calls if c == ("POST", "/vm/deallocate")) == mod.DEALLOC_ATTEMPTS bound = 3600 + mod.DEALLOC_ATTEMPTS * (mod.DEALLOC_POLL_S + mod.CONFIRM_S + 2 * mod.POLL_EVERY_S) + 300 assert t[0] < bound # bounded even when Azure never confirms def test_unreadable_power_state_never_counts_as_deallocated(orch): mod, t, args = orch mod.az = FakeAzure(t, run_ends_after=600) real_arm = mod.az.arm def arm(method, path, body=None, api=None): if path.endswith("/instanceView") and mod.az.power == "VM deallocated" and ("POST", "/vm/deallocate") in mod.az.calls: raise ConnectionError("no answer") return real_arm(method, path, body, api) mod.az.arm = arm assert mod.main(args) == 3 def test_no_credential_route_exists(): src = open(os.path.join(HERE, "probe_orchestrator.py"), encoding="utf-8").read() tpl = open(os.path.join(HERE, "vm_probe_session.tpl.ps1"), encoding="utf-8").read() for word in ("OXU", "OXP", "handover", "cred_send", "RSACryptoServiceProvider"): assert word not in src and word not in tpl, word assert not os.path.exists(os.path.join(HERE, "cred_send.ps1"))