From e7e8c7fdb4c6edcb0cd19560db700c6dcf16527c Mon Sep 17 00:00:00 2001 From: claude Date: Mon, 6 Jul 2026 17:47:52 +0000 Subject: [PATCH] refactor: extract supervise types to _supervise_types, eliminate get_supervise_mod() Proposal, Response, AuditEntry, host_db_path, bot_bottle_root, and their shared constants/helpers are moved to a new _supervise_types.py module. queue_store and audit_store now import these directly instead of deferring through get_supervise_mod() at call time. host_db_path() in _supervise_types uses sys.modules[__name__].bot_bottle_root() so monkey-patches on _supervise_types.bot_bottle_root propagate through it. supervise.host_db_path() is re-defined to call bot_bottle_root() through supervise's own globals, preserving patchability for callers that go through the supervise module. Test fixtures that create stores without an explicit db_path now patch _sv_types.bot_bottle_root alongside supervise.bot_bottle_root. The server test patches both flat and package _supervise_types since handle_tools_call runs in the package context while the responder thread uses the flat module. Dockerfile.sidecars gets a COPY for _supervise_types.py so the sidecar bundle retains its flat import chain. Follow-up to issuecomment-2872 on PR #320. --- Dockerfile.sidecars | 1 + bot_bottle/_supervise_types.py | 182 +++++++++++++++++++++++++ bot_bottle/audit_store.py | 28 +--- bot_bottle/queue_store.py | 33 +---- bot_bottle/supervise.py | 199 ++++++---------------------- tests/unit/test_supervise.py | 34 ++++- tests/unit/test_supervise_cli.py | 12 +- tests/unit/test_supervise_server.py | 16 ++- 8 files changed, 288 insertions(+), 217 deletions(-) create mode 100644 bot_bottle/_supervise_types.py diff --git a/Dockerfile.sidecars b/Dockerfile.sidecars index 285785a..3feaf4d 100644 --- a/Dockerfile.sidecars +++ b/Dockerfile.sidecars @@ -68,6 +68,7 @@ COPY bot_bottle/dlp_detectors.py /app/dlp_detectors.py COPY bot_bottle/yaml_subset.py /app/yaml_subset.py COPY bot_bottle/migrations.py /app/migrations.py COPY bot_bottle/db_store.py /app/db_store.py +COPY bot_bottle/_supervise_types.py /app/_supervise_types.py COPY bot_bottle/queue_store.py /app/queue_store.py COPY bot_bottle/audit_store.py /app/audit_store.py COPY bot_bottle/store_manager.py /app/store_manager.py diff --git a/bot_bottle/_supervise_types.py b/bot_bottle/_supervise_types.py new file mode 100644 index 0000000..d71cdfe --- /dev/null +++ b/bot_bottle/_supervise_types.py @@ -0,0 +1,182 @@ +"""Shared types and path helpers for supervise stores (PRD 0013). + +Extracted from supervise.py so queue_store and audit_store can import +Proposal, Response, AuditEntry, and host_db_path without creating a +circular import (supervise imports from queue_store/audit_store and +vice-versa). + +Patching bot_bottle_root on this module propagates through host_db_path +because host_db_path looks up bot_bottle_root via sys.modules[__name__] +at call time rather than capturing it at import time. +""" + +from __future__ import annotations + +import dataclasses +import sys +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + + +HOST_DB_FILENAME = "bot-bottle.db" + +TOOL_EGRESS_BLOCK = "egress-block" +TOOL_EGRESS_ALLOW = "egress-allow" +TOOL_GITLEAKS_ALLOW = "gitleaks-allow" +TOOL_EGRESS_TOKEN_ALLOW = "egress-token-allow" +TOOL_LIST_EGRESS_ROUTES = "list-egress-routes" +TOOLS: tuple[str, ...] = ( + TOOL_EGRESS_ALLOW, + TOOL_EGRESS_BLOCK, + TOOL_GITLEAKS_ALLOW, + TOOL_EGRESS_TOKEN_ALLOW, + TOOL_LIST_EGRESS_ROUTES, +) + +STATUS_APPROVED = "approved" +STATUS_MODIFIED = "modified" +STATUS_REJECTED = "rejected" +STATUSES: tuple[str, ...] = (STATUS_APPROVED, STATUS_MODIFIED, STATUS_REJECTED) + +ACTION_OPERATOR_EDIT = "operator-edit" + + +def bot_bottle_root() -> Path: + return Path.home() / ".bot-bottle" + + +def host_db_path() -> Path: + # Look up bot_bottle_root through this module's live namespace so that + # monkey-patches on _supervise_types.bot_bottle_root take effect. + return sys.modules[__name__].bot_bottle_root() / HOST_DB_FILENAME + + +def _require_str(raw: dict[str, object], key: str) -> str: + value = raw.get(key) + if not isinstance(value, str): + raise ValueError(f"missing or non-string field {key!r}") + return value + + +@dataclass(frozen=True) +class Proposal: + """One pending tool-call from the agent.""" + + id: str + bottle_slug: str + tool: str + proposed_file: str + justification: str + arrival_timestamp: str + current_file_hash: str + + @classmethod + def new( + cls, + *, + bottle_slug: str, + tool: str, + proposed_file: str, + justification: str, + current_file_hash: str, + now: datetime | None = None, + ) -> "Proposal": + ts = (now or datetime.now(timezone.utc)).isoformat() + return cls( + id=str(uuid.uuid4()), + bottle_slug=bottle_slug, + tool=tool, + proposed_file=proposed_file, + justification=justification, + arrival_timestamp=ts, + current_file_hash=current_file_hash, + ) + + def to_dict(self) -> dict[str, object]: + return dataclasses.asdict(self) + + @classmethod + def from_dict(cls, raw: dict[str, object]) -> "Proposal": + tool = _require_str(raw, "tool") + if tool not in TOOLS: + raise ValueError(f"tool must be one of {TOOLS}; got {tool!r}") + return cls( + id=_require_str(raw, "id"), + bottle_slug=_require_str(raw, "bottle_slug"), + tool=tool, + proposed_file=_require_str(raw, "proposed_file"), + justification=_require_str(raw, "justification"), + arrival_timestamp=_require_str(raw, "arrival_timestamp"), + current_file_hash=_require_str(raw, "current_file_hash"), + ) + + +@dataclass(frozen=True) +class Response: + """The operator's decision on a proposal.""" + + proposal_id: str + status: str + notes: str + final_file: str | None = None + + def to_dict(self) -> dict[str, object]: + return dataclasses.asdict(self) + + @classmethod + def from_dict(cls, raw: dict[str, object]) -> "Response": + status = _require_str(raw, "status") + if status not in STATUSES: + raise ValueError( + f"response status must be one of {STATUSES}; got {status!r}" + ) + final = raw.get("final_file") + if final is not None and not isinstance(final, str): + raise ValueError( + f"final_file must be a string or null; got {type(final).__name__}" + ) + return cls( + proposal_id=_require_str(raw, "proposal_id"), + status=status, + notes=_require_str(raw, "notes"), + final_file=final, + ) + + +@dataclass(frozen=True) +class AuditEntry: + """One row of the per-bottle audit log.""" + + timestamp: str + bottle_slug: str + component: str + operator_action: str + operator_notes: str + justification: str + diff: str + + def to_dict(self) -> dict[str, object]: + return dataclasses.asdict(self) + + +__all__ = [ + "ACTION_OPERATOR_EDIT", + "AuditEntry", + "HOST_DB_FILENAME", + "Proposal", + "Response", + "STATUSES", + "STATUS_APPROVED", + "STATUS_MODIFIED", + "STATUS_REJECTED", + "TOOLS", + "TOOL_EGRESS_ALLOW", + "TOOL_EGRESS_BLOCK", + "TOOL_EGRESS_TOKEN_ALLOW", + "TOOL_GITLEAKS_ALLOW", + "TOOL_LIST_EGRESS_ROUTES", + "bot_bottle_root", + "host_db_path", +] diff --git a/bot_bottle/audit_store.py b/bot_bottle/audit_store.py index 5f848d0..7411a10 100644 --- a/bot_bottle/audit_store.py +++ b/bot_bottle/audit_store.py @@ -4,35 +4,17 @@ from __future__ import annotations import sqlite3 from pathlib import Path -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from .supervise import AuditEntry try: + from ._supervise_types import AuditEntry, host_db_path from .db_store import DbStore from .migrations import TableMigrations except ImportError: + from _supervise_types import AuditEntry, host_db_path # type: ignore[import-not-found] # pylint: disable=import-error,no-name-in-module from db_store import DbStore # type: ignore[import-not-found] # pylint: disable=import-error,no-name-in-module from migrations import TableMigrations # type: ignore[import-not-found] # pylint: disable=import-error,no-name-in-module -def get_supervise_mod() -> object: - """Lazy import of supervise to avoid a circular-import at module init time. - Mirrors our own module identity so patches on supervise.bot_bottle_root - propagate correctly in both flat (sidecar / sys.path-injection tests) and - package contexts.""" - import sys - sv_name = "supervise" if __name__ == "audit_store" else "bot_bottle.supervise" - if sv_name in sys.modules: - return sys.modules[sv_name] - try: - import bot_bottle.supervise as _m - except ImportError: - import supervise as _m # type: ignore[import-not-found] # pylint: disable=import-error,no-name-in-module - return _m - - # One entry per schema version: _MIGRATIONS.migrations[0] brings a fresh DB # to version 1, [1] to version 2, and so on. Add new migrations at the end; # never edit existing ones. @@ -57,8 +39,7 @@ class AuditStore(DbStore): """SQLite-backed persistent store for supervise audit entries.""" def __init__(self, db_path: Path | None = None) -> None: - resolved = db_path or get_supervise_mod().host_db_path() # type: ignore[attr-defined] - super().__init__(resolved, _MIGRATIONS) + super().__init__(db_path or host_db_path(), _MIGRATIONS) def write_audit_entry(self, entry: AuditEntry) -> Path: with self._connect() as conn: @@ -98,8 +79,7 @@ class AuditStore(DbStore): @staticmethod def _row_to_entry(row: sqlite3.Row) -> AuditEntry: - m = get_supervise_mod() - return m.AuditEntry( # type: ignore[attr-defined] + return AuditEntry( timestamp=row["timestamp"], bottle_slug=row["bottle_slug"], component=row["component"], diff --git a/bot_bottle/queue_store.py b/bot_bottle/queue_store.py index 76a02da..17a0a0b 100644 --- a/bot_bottle/queue_store.py +++ b/bot_bottle/queue_store.py @@ -5,38 +5,17 @@ from __future__ import annotations import os import sqlite3 from pathlib import Path -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - from .supervise import Proposal, Response try: + from ._supervise_types import Proposal, Response, host_db_path from .db_store import DbStore from .migrations import TableMigrations except ImportError: + from _supervise_types import Proposal, Response, host_db_path # type: ignore[import-not-found] # pylint: disable=import-error,no-name-in-module from db_store import DbStore # type: ignore[import-not-found] # pylint: disable=import-error,no-name-in-module from migrations import TableMigrations # type: ignore[import-not-found] # pylint: disable=import-error,no-name-in-module -def get_supervise_mod() -> object: - """Lazy import of supervise to avoid a circular-import at module init time. - By the time any QueueStore method is called, both modules are fully loaded. - - Mirrors our own module identity: when we are 'queue_store' (sidecar flat - context or tests that inject bot_bottle/ into sys.path) we use the flat - 'supervise' module so that patches on supervise.bot_bottle_root propagate - correctly. When we are 'bot_bottle.queue_store' we use 'bot_bottle.supervise'.""" - import sys - sv_name = "supervise" if __name__ == "queue_store" else "bot_bottle.supervise" - if sv_name in sys.modules: - return sys.modules[sv_name] - try: - import bot_bottle.supervise as _m - except ImportError: - import supervise as _m # type: ignore[import-not-found] # pylint: disable=import-error,no-name-in-module - return _m - - # One entry per schema version: _MIGRATIONS.migrations[0] brings a fresh DB # to version 1, [1] to version 2, and so on. Add new migrations at the end; # never edit existing ones. @@ -83,7 +62,7 @@ class QueueStore(DbStore): # bind-mounted host DB. On the host this env var is never set, # so we always fall through to host_db_path(). env_path = os.environ.get("SUPERVISE_DB_PATH", "").strip() - resolved = Path(env_path) if env_path else get_supervise_mod().host_db_path() # type: ignore[attr-defined] + resolved = Path(env_path) if env_path else host_db_path() super().__init__(resolved, _MIGRATIONS) def write_proposal(self, proposal: Proposal) -> Path: @@ -215,8 +194,7 @@ class QueueStore(DbStore): @staticmethod def _row_to_proposal(row: sqlite3.Row) -> Proposal: - m = get_supervise_mod() - return m.Proposal( # type: ignore[attr-defined] + return Proposal( id=row["id"], bottle_slug=row["bottle_slug"], tool=row["tool"], @@ -228,8 +206,7 @@ class QueueStore(DbStore): @staticmethod def _row_to_response(row: sqlite3.Row) -> Response: - m = get_supervise_mod() - return m.Response( # type: ignore[attr-defined] + return Response( proposal_id=row["proposal_id"], status=row["status"], notes=row["notes"], diff --git a/bot_bottle/supervise.py b/bot_bottle/supervise.py index 0a9fcad..5dd1dc5 100644 --- a/bot_bottle/supervise.py +++ b/bot_bottle/supervise.py @@ -30,35 +30,56 @@ remediation engines that wire real config changes land in PRDs 0014, from __future__ import annotations -import dataclasses import difflib import hashlib import time -import uuid from abc import ABC from dataclasses import dataclass -from datetime import datetime, timezone from pathlib import Path +try: + from ._supervise_types import ( + ACTION_OPERATOR_EDIT, + AuditEntry, + HOST_DB_FILENAME, + Proposal, + Response, + STATUSES, + STATUS_APPROVED, + STATUS_MODIFIED, + STATUS_REJECTED, + TOOLS, + TOOL_EGRESS_ALLOW, + TOOL_EGRESS_BLOCK, + TOOL_EGRESS_TOKEN_ALLOW, + TOOL_GITLEAKS_ALLOW, + TOOL_LIST_EGRESS_ROUTES, + bot_bottle_root, + ) +except ImportError: + from _supervise_types import ( # type: ignore[import-not-found,no-redef] # pylint: disable=import-error,no-name-in-module + ACTION_OPERATOR_EDIT, + AuditEntry, + HOST_DB_FILENAME, + Proposal, + Response, + STATUSES, + STATUS_APPROVED, + STATUS_MODIFIED, + STATUS_REJECTED, + TOOLS, + TOOL_EGRESS_ALLOW, + TOOL_EGRESS_BLOCK, + TOOL_EGRESS_TOKEN_ALLOW, + TOOL_GITLEAKS_ALLOW, + TOOL_LIST_EGRESS_ROUTES, + bot_bottle_root, + ) + SUPERVISE_HOSTNAME = "supervise" SUPERVISE_PORT = 9100 -TOOL_EGRESS_BLOCK = "egress-block" -TOOL_EGRESS_ALLOW = "egress-allow" -TOOL_GITLEAKS_ALLOW = "gitleaks-allow" -# Written directly by the egress addon (not an agent-facing MCP tool) when an -# outbound DLP token block is routed to the operator for override (PRD 0062). -TOOL_EGRESS_TOKEN_ALLOW = "egress-token-allow" -TOOL_LIST_EGRESS_ROUTES = "list-egress-routes" -TOOLS: tuple[str, ...] = ( - TOOL_EGRESS_ALLOW, - TOOL_EGRESS_BLOCK, - TOOL_GITLEAKS_ALLOW, - TOOL_EGRESS_TOKEN_ALLOW, - TOOL_LIST_EGRESS_ROUTES, -) - # The supervise sidecar uses these to query egress's # introspection endpoint for the `list-egress-routes` MCP # tool. The hostname + port match egress's docker network @@ -74,149 +95,26 @@ COMPONENT_FOR_TOOL: dict[str, str] = { TOOL_EGRESS_BLOCK: "egress", } -STATUS_APPROVED = "approved" -STATUS_MODIFIED = "modified" -STATUS_REJECTED = "rejected" -STATUSES: tuple[str, ...] = (STATUS_APPROVED, STATUS_MODIFIED, STATUS_REJECTED) - -# Operator-initiated audit entries (no tool call). PRD 0014's -# `routes edit ` verb writes entries with this action. -ACTION_OPERATOR_EDIT = "operator-edit" - DB_PATH_IN_CONTAINER = "/run/supervise/bot-bottle.db" DEFAULT_POLL_INTERVAL_SEC = 0.5 -HOST_DB_FILENAME = "bot-bottle.db" # --- Paths ----------------------------------------------------------------- -def bot_bottle_root() -> Path: - return Path.home() / ".bot-bottle" - - def audit_dir() -> Path: return bot_bottle_root() / "audit" -def audit_log_path(component: str, slug: str) -> Path: - return audit_dir() / f"{component}-{slug}.log" - - def host_db_path() -> Path: + # Calls bot_bottle_root() through this module's globals so that patches + # on supervise.bot_bottle_root propagate to callers going through + # supervise.host_db_path(). return bot_bottle_root() / HOST_DB_FILENAME -# --- Dataclasses ----------------------------------------------------------- - - -@dataclass(frozen=True) -class Proposal: - """One pending tool-call from the agent.""" - - id: str - bottle_slug: str - tool: str - proposed_file: str - justification: str - arrival_timestamp: str - current_file_hash: str - - @classmethod - def new( - cls, - *, - bottle_slug: str, - tool: str, - proposed_file: str, - justification: str, - current_file_hash: str, - now: datetime | None = None, - ) -> "Proposal": - ts = (now or datetime.now(timezone.utc)).isoformat() - return cls( - id=str(uuid.uuid4()), - bottle_slug=bottle_slug, - tool=tool, - proposed_file=proposed_file, - justification=justification, - arrival_timestamp=ts, - current_file_hash=current_file_hash, - ) - - def to_dict(self) -> dict[str, object]: - return dataclasses.asdict(self) - - @classmethod - def from_dict(cls, raw: dict[str, object]) -> "Proposal": - tool = _require_str(raw, "tool") - if tool not in TOOLS: - raise ValueError(f"tool must be one of {TOOLS}; got {tool!r}") - return cls( - id=_require_str(raw, "id"), - bottle_slug=_require_str(raw, "bottle_slug"), - tool=tool, - proposed_file=_require_str(raw, "proposed_file"), - justification=_require_str(raw, "justification"), - arrival_timestamp=_require_str(raw, "arrival_timestamp"), - current_file_hash=_require_str(raw, "current_file_hash"), - ) - - -@dataclass(frozen=True) -class Response: - """The operator's decision on a proposal. The TUI writes one of - these to the queue table; the sidecar reads it and returns the - `{status, notes}` pair to the agent's tool call. - - `final_file` carries the file content the supervisor will - actually apply: for `approved`, equal to the proposal's - `proposed_file`; for `modified`, the operator's edited version - (the audit diff is current → final_file, not current → - proposed_file); for `rejected`, None.""" - - proposal_id: str - status: str - notes: str - final_file: str | None = None - - def to_dict(self) -> dict[str, object]: - return dataclasses.asdict(self) - - @classmethod - def from_dict(cls, raw: dict[str, object]) -> "Response": - status = _require_str(raw, "status") - if status not in STATUSES: - raise ValueError( - f"response status must be one of {STATUSES}; got {status!r}" - ) - final = raw.get("final_file") - if final is not None and not isinstance(final, str): - raise ValueError( - f"final_file must be a string or null; got {type(final).__name__}" - ) - return cls( - proposal_id=_require_str(raw, "proposal_id"), - status=status, - notes=_require_str(raw, "notes"), - final_file=final, - ) - - -@dataclass(frozen=True) -class AuditEntry: - """One row of the per-bottle audit log. JSON-Lines, append-only.""" - - timestamp: str - bottle_slug: str - component: str - operator_action: str - operator_notes: str - justification: str - diff: str - - def to_dict(self) -> dict[str, object]: - return dataclasses.asdict(self) +def audit_log_path(component: str, slug: str) -> Path: + return audit_dir() / f"{component}-{slug}.log" try: @@ -365,15 +263,6 @@ class Supervise(ABC): db_path=mgr.db_path, ) -# --- Helpers --------------------------------------------------------------- - - -def _require_str(raw: dict[str, object], key: str) -> str: - value = raw.get(key) - if not isinstance(value, str): - raise ValueError(f"missing or non-string field {key!r}") - return value - __all__ = [ "ACTION_OPERATOR_EDIT", diff --git a/tests/unit/test_supervise.py b/tests/unit/test_supervise.py index 348a69b..a75c19f 100644 --- a/tests/unit/test_supervise.py +++ b/tests/unit/test_supervise.py @@ -8,6 +8,7 @@ from datetime import datetime, timezone from pathlib import Path from bot_bottle import supervise +from bot_bottle import _supervise_types as _sv_types from bot_bottle.audit_store import AuditStore from bot_bottle.queue_store import QueueStore from bot_bottle.supervise import ( @@ -122,13 +123,20 @@ class TestQueueIO(unittest.TestCase): self._tmp.cleanup() def _patch_home(self, fake_home: Path): - original = supervise.bot_bottle_root + original_sv = supervise.bot_bottle_root + original_svt = _sv_types.bot_bottle_root def fake_root() -> Path: return fake_home / ".bot-bottle" supervise.bot_bottle_root = fake_root # type: ignore[assignment] - return lambda: setattr(supervise, "bot_bottle_root", original) + _sv_types.bot_bottle_root = fake_root # type: ignore[assignment] + + def restore() -> None: + supervise.bot_bottle_root = original_sv # type: ignore[assignment] + _sv_types.bot_bottle_root = original_svt # type: ignore[assignment] + + return restore def test_write_and_read_proposal(self): p = _proposal() @@ -232,13 +240,20 @@ class TestAuditLog(unittest.TestCase): self._tmp.cleanup() def _patch_home(self, fake_home: Path): - original = supervise.bot_bottle_root + original_sv = supervise.bot_bottle_root + original_svt = _sv_types.bot_bottle_root def fake_root() -> Path: return fake_home / ".bot-bottle" supervise.bot_bottle_root = fake_root # type: ignore[assignment] - return lambda: setattr(supervise, "bot_bottle_root", original) + _sv_types.bot_bottle_root = fake_root # type: ignore[assignment] + + def restore() -> None: + supervise.bot_bottle_root = original_sv # type: ignore[assignment] + _sv_types.bot_bottle_root = original_svt # type: ignore[assignment] + + return restore def test_write_then_read_single_entry(self): e = AuditEntry( @@ -385,13 +400,20 @@ class TestSupervisePrepare(unittest.TestCase): self._tmp.cleanup() def _patch_home(self, fake_home: Path): - original = supervise.bot_bottle_root + original_sv = supervise.bot_bottle_root + original_svt = _sv_types.bot_bottle_root def fake_root() -> Path: return fake_home / ".bot-bottle" supervise.bot_bottle_root = fake_root # type: ignore[assignment] - return lambda: setattr(supervise, "bot_bottle_root", original) + _sv_types.bot_bottle_root = fake_root # type: ignore[assignment] + + def restore() -> None: + supervise.bot_bottle_root = original_sv # type: ignore[assignment] + _sv_types.bot_bottle_root = original_svt # type: ignore[assignment] + + return restore def test_prepare_creates_queue(self): plan = _StubSupervise().prepare("dev", self.stage_dir) diff --git a/tests/unit/test_supervise_cli.py b/tests/unit/test_supervise_cli.py index bcf10a0..9d1dc66 100644 --- a/tests/unit/test_supervise_cli.py +++ b/tests/unit/test_supervise_cli.py @@ -12,6 +12,7 @@ from pathlib import Path from unittest.mock import patch from bot_bottle import supervise +from bot_bottle import _supervise_types as _sv_types from bot_bottle.audit_store import AuditStore from bot_bottle.cli import supervise as supervise_cli from bot_bottle.queue_store import QueueStore @@ -54,13 +55,20 @@ class _FakeHomeMixin: def _setup_fake_home(self): self._tmp = tempfile.TemporaryDirectory(prefix="supervise-test.") - original = supervise.bot_bottle_root + original_sv = supervise.bot_bottle_root + original_svt = _sv_types.bot_bottle_root def fake_root() -> Path: return Path(self._tmp.name) / ".bot-bottle" supervise.bot_bottle_root = fake_root # type: ignore[assignment] - self._restore_home = lambda: setattr(supervise, "bot_bottle_root", original) + _sv_types.bot_bottle_root = fake_root # type: ignore[assignment] + + def restore() -> None: + supervise.bot_bottle_root = original_sv # type: ignore[assignment] + _sv_types.bot_bottle_root = original_svt # type: ignore[assignment] + + self._restore_home = restore QueueStore("").migrate() AuditStore().migrate() diff --git a/tests/unit/test_supervise_server.py b/tests/unit/test_supervise_server.py index 51db278..ca8573d 100644 --- a/tests/unit/test_supervise_server.py +++ b/tests/unit/test_supervise_server.py @@ -19,8 +19,10 @@ sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent / "bot_bott import supervise as _sv # noqa: E402 # type: ignore import queue_store as _qs # noqa: E402 # type: ignore import audit_store as _as # noqa: E402 # type: ignore +import _supervise_types as _svt_flat # noqa: E402 # type: ignore from bot_bottle import supervise_server # noqa: E402 +from bot_bottle import _supervise_types as _svt_pkg # noqa: E402 from bot_bottle.supervise_server import ( ERR_INTERNAL, ERR_INVALID_PARAMS, @@ -277,13 +279,23 @@ class TestHandleToolsCall(unittest.TestCase): self._tmp.cleanup() def _patch_home(self, fake_home: Path): - original = _sv.bot_bottle_root + original_sv = _sv.bot_bottle_root + original_flat = _svt_flat.bot_bottle_root + original_pkg = _svt_pkg.bot_bottle_root def fake_root() -> Path: return fake_home / ".bot-bottle" _sv.bot_bottle_root = fake_root # type: ignore[assignment] - return lambda: setattr(_sv, "bot_bottle_root", original) + _svt_flat.bot_bottle_root = fake_root # type: ignore[assignment] + _svt_pkg.bot_bottle_root = fake_root # type: ignore[assignment] + + def restore() -> None: + _sv.bot_bottle_root = original_sv # type: ignore[assignment] + _svt_flat.bot_bottle_root = original_flat # type: ignore[assignment] + _svt_pkg.bot_bottle_root = original_pkg # type: ignore[assignment] + + return restore def _respond_when_proposal_appears(self, status: str, notes: str = "") -> threading.Thread: """Background thread: poll the queue for a fresh proposal, write a