test(api): bring two command test modules under strict pyrefly (#45633)

Co-authored-by: yetuge <yetuge@users.noreply.github.com>

api/tests/unit_tests/commands/test_retention.py

@@ -15,7 +15,6 @@ from sqlalchemy.orm import Session, scoped_session, sessionmaker
 
 from commands import clean_expired_messages, retention
 from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus
-from models.base import TypeBase
 from models.enums import CreatorUserRole, WorkflowRunTriggeredFrom
 from models.workflow import (
     WorkflowNodeExecutionModel,
@@ -165,7 +164,7 @@ class TestArchiveWorkflowRunsCases:
         "command",
         [retention.restore_workflow_runs, retention.delete_archived_workflow_runs],
     )
-    def test_v2_archive_maintenance_rejects_explicitly_empty_tenant_ids(command):
+    def test_v2_archive_maintenance_rejects_explicitly_empty_tenant_ids(command: click.Command) -> None:
         result = CliRunner().invoke(
             command,
             ["--tenant-ids", "", "--target-month", "2025-03"],
@@ -175,20 +174,16 @@ class TestArchiveWorkflowRunsCases:
         assert "tenant-ids must not be empty" in result.output
 
     @staticmethod
-    def test_archive_tenant_id_parser_keeps_omitted_scope_unset():
+    def test_archive_tenant_id_parser_keeps_omitted_scope_unset() -> None:
         assert retention._parse_comma_separated_ids(None, param_name="tenant-ids") is None
 
     @staticmethod
     @pytest.fixture
-    def archive_db(sqlite_engine: Engine) -> ArchiveDatabase:
-        """Create only the workflow tables used by archive planning."""
+    def archive_db(sqlite_session: Session) -> ArchiveDatabase:
+        """Expose archive candidates through the shared full-schema SQLite database."""
 
-        TypeBase.metadata.create_all(
-            sqlite_engine,
-            tables=[WorkflowRun.__table__, WorkflowNodeExecutionModel.__table__],
-        )
         return ArchiveDatabase(
-            session_maker=sessionmaker(bind=sqlite_engine, expire_on_commit=False),
+            session_maker=sessionmaker(bind=sqlite_session.get_bind(), expire_on_commit=False),
             end_before=datetime.datetime(2025, 4, 1, tzinfo=datetime.UTC),
         )
 
@@ -349,8 +344,10 @@ class TestArchiveWorkflowRunsCases:
             Mock(side_effect=RuntimeError("tenant plan failed")),
         )
 
+        callback = retention.archive_workflow_runs.callback
+        assert callback is not None
         with pytest.raises(click.ClickException, match="Failed to resolve workflow archive tenant plan"):
-            retention.archive_workflow_runs.callback(
+            callback(
                 tenant_ids="tenant-1",
                 tenant_prefixes=None,
                 before_days=90,
@@ -370,7 +367,7 @@ class TestArchiveWorkflowRunsCases:
     @staticmethod
     def test_delete_archived_workflow_runs_keeps_single_page_behavior_without_all_pages(
         monkeypatch: pytest.MonkeyPatch,
-    ):
+    ) -> None:
         deleter = _patch_bundle_deleter(
             monkeypatch,
             [_delete_summary(processed=2, succeeded=2, next_catalog_id=_CURSOR_1)],
@@ -389,7 +386,9 @@ class TestArchiveWorkflowRunsCases:
         assert deleter.delete_batch.call_args.kwargs["limit"] == 2
 
     @staticmethod
-    def test_delete_archived_workflow_runs_all_pages_continues_until_empty_page(monkeypatch: pytest.MonkeyPatch):
+    def test_delete_archived_workflow_runs_all_pages_continues_until_empty_page(
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
         deleter = _patch_bundle_deleter(
             monkeypatch,
             [
@@ -414,7 +413,7 @@ class TestArchiveWorkflowRunsCases:
     @staticmethod
     def test_delete_archived_workflow_runs_all_pages_fetches_empty_page_after_exact_full_page(
         monkeypatch: pytest.MonkeyPatch,
-    ):
+    ) -> None:
         deleter = _patch_bundle_deleter(
             monkeypatch,
             [
@@ -433,7 +432,9 @@ class TestArchiveWorkflowRunsCases:
         assert deleter.delete_batch.call_args_list[1].kwargs["after_catalog_id"] == _CURSOR_1
 
     @staticmethod
-    def test_delete_archived_workflow_runs_all_pages_stops_at_first_failed_page(monkeypatch: pytest.MonkeyPatch):
+    def test_delete_archived_workflow_runs_all_pages_stops_at_first_failed_page(
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
         failed_result = BundleOperationResult(
             catalog_id=_CURSOR_2,
             bundle_id="bundle-failed",
@@ -464,7 +465,7 @@ class TestArchiveWorkflowRunsCases:
     @staticmethod
     def test_delete_archived_workflow_runs_all_pages_fails_when_cursor_does_not_advance(
         monkeypatch: pytest.MonkeyPatch,
-    ):
+    ) -> None:
         deleter = _patch_bundle_deleter(
             monkeypatch,
             [_delete_summary(processed=1, succeeded=1, next_catalog_id=None)],
@@ -480,7 +481,9 @@ class TestArchiveWorkflowRunsCases:
         assert "cursor did not advance" in result.output.lower()
 
     @staticmethod
-    def test_delete_archived_workflow_runs_all_pages_uses_preview_cursor_for_dry_run(monkeypatch: pytest.MonkeyPatch):
+    def test_delete_archived_workflow_runs_all_pages_uses_preview_cursor_for_dry_run(
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
         deleter = _patch_bundle_deleter(
             monkeypatch,
             [
@@ -503,7 +506,7 @@ class TestArchiveWorkflowRunsCases:
     @staticmethod
     def test_delete_archived_workflow_runs_dry_run_failure_separates_preview_and_destructive_cursors(
         monkeypatch: pytest.MonkeyPatch,
-    ):
+    ) -> None:
         failed_result = BundleOperationResult(
             catalog_id=_CURSOR_2,
             bundle_id="bundle-failed",
@@ -546,7 +549,9 @@ class TestArchiveWorkflowRunsCases:
         assert f"destructive_retry_after_catalog_id={_CURSOR_0}" in result.output
 
     @staticmethod
-    def test_delete_archived_workflow_runs_all_pages_starts_after_explicit_cursor(monkeypatch: pytest.MonkeyPatch):
+    def test_delete_archived_workflow_runs_all_pages_starts_after_explicit_cursor(
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
         deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
 
         result = CliRunner().invoke(
@@ -576,7 +581,9 @@ class TestArchiveWorkflowRunsCases:
             ["--run-shard-index", "0", "--run-shard-total", "17"],
         ],
     )
-    def test_delete_archived_workflow_runs_rejects_invalid_run_shard_options(monkeypatch, shard_args):
+    def test_delete_archived_workflow_runs_rejects_invalid_run_shard_options(
+        monkeypatch: pytest.MonkeyPatch, shard_args: list[str]
+    ) -> None:
         deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
 
         result = CliRunner().invoke(
@@ -588,7 +595,9 @@ class TestArchiveWorkflowRunsCases:
         deleter.delete_batch.assert_not_called()
 
     @staticmethod
-    def test_delete_archived_workflow_runs_passes_formatted_run_shard_to_service(monkeypatch: pytest.MonkeyPatch):
+    def test_delete_archived_workflow_runs_passes_formatted_run_shard_to_service(
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
         deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
 
         result = CliRunner().invoke(
@@ -616,7 +625,9 @@ class TestArchiveWorkflowRunsCases:
         assert deleter.delete_batch.call_args.kwargs["shard"] == "03-of-16"
 
     @staticmethod
-    def test_delete_archived_workflow_runs_rejects_mixed_catalog_shards_before_delete(monkeypatch: pytest.MonkeyPatch):
+    def test_delete_archived_workflow_runs_rejects_mixed_catalog_shards_before_delete(
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
         deleter = _patch_bundle_deleter(monkeypatch, [_delete_summary(processed=0)])
         deleter.validate_catalog_shards.side_effect = ValueError("unexpected shards: 00-of-01")
 

api/tests/unit_tests/commands/test_system.py

@@ -2,6 +2,7 @@ import sys
 import threading
 import types
 import uuid
+from collections.abc import Callable
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
@@ -318,26 +319,37 @@ def test_reset_iterates_all_tenants(monkeypatch: pytest.MonkeyPatch, sqlite_sess
 HEARTBEAT_WAIT_TIMEOUT_SECONDS = 5.0
 
 
-def _install_fake_flask_migrate(monkeypatch, upgrade_impl) -> None:
+def _install_fake_flask_migrate(monkeypatch: pytest.MonkeyPatch, upgrade_impl: Callable[[], None]) -> None:
     module = types.ModuleType("flask_migrate")
-    module.upgrade = upgrade_impl
+    # ModuleType does not declare `upgrade`; write it through the namespace mapping.
+    module.__dict__["upgrade"] = upgrade_impl
     monkeypatch.setitem(sys.modules, "flask_migrate", module)
 
 
+def _redis_client_mock() -> MagicMock:
+    client = system_commands.redis_client
+    assert isinstance(client, MagicMock)
+    return client
+
+
 def _invoke_upgrade_db() -> int:
+    callback = commands.upgrade_db.callback
+    assert callback is not None
     try:
-        commands.upgrade_db.callback()
+        callback()
     except SystemExit as e:
         return int(e.code or 0)
     return 0
 
 
-def test_upgrade_db_skips_when_lock_not_acquired(monkeypatch, capsys):
+def test_upgrade_db_skips_when_lock_not_acquired(
+    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
+) -> None:
     monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 1234)
 
     lock = MagicMock()
     lock.acquire.return_value = False
-    system_commands.redis_client.lock.return_value = lock
+    _redis_client_mock().lock.return_value = lock
 
     exit_code = _invoke_upgrade_db()
     captured = capsys.readouterr()
@@ -345,20 +357,22 @@ def test_upgrade_db_skips_when_lock_not_acquired(monkeypatch, capsys):
     assert exit_code == 0
     assert "Database migration skipped" in captured.out
 
-    system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=1234, thread_local=False)
+    _redis_client_mock().lock.assert_called_once_with(name="db_upgrade_lock", timeout=1234, thread_local=False)
     lock.acquire.assert_called_once_with(blocking=False)
     lock.release.assert_not_called()
 
 
-def test_upgrade_db_failure_not_masked_by_lock_release(monkeypatch, capsys):
+def test_upgrade_db_failure_not_masked_by_lock_release(
+    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
+) -> None:
     monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 321)
 
     lock = MagicMock()
     lock.acquire.return_value = True
     lock.release.side_effect = LockNotOwnedError("simulated")
-    system_commands.redis_client.lock.return_value = lock
+    _redis_client_mock().lock.return_value = lock
 
-    def _upgrade():
+    def _upgrade() -> None:
         raise RuntimeError("boom")
 
     _install_fake_flask_migrate(monkeypatch, _upgrade)
@@ -369,18 +383,20 @@ def test_upgrade_db_failure_not_masked_by_lock_release(monkeypatch, capsys):
     assert exit_code == 1
     assert "Database migration failed: boom" in captured.out
 
-    system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=321, thread_local=False)
+    _redis_client_mock().lock.assert_called_once_with(name="db_upgrade_lock", timeout=321, thread_local=False)
     lock.acquire.assert_called_once_with(blocking=False)
     lock.release.assert_called_once()
 
 
-def test_upgrade_db_success_ignores_lock_not_owned_on_release(monkeypatch, capsys):
+def test_upgrade_db_success_ignores_lock_not_owned_on_release(
+    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
+) -> None:
     monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 999)
 
     lock = MagicMock()
     lock.acquire.return_value = True
     lock.release.side_effect = LockNotOwnedError("simulated")
-    system_commands.redis_client.lock.return_value = lock
+    _redis_client_mock().lock.return_value = lock
 
     _install_fake_flask_migrate(monkeypatch, lambda: None)
 
@@ -390,12 +406,14 @@ def test_upgrade_db_success_ignores_lock_not_owned_on_release(monkeypatch, capsy
     assert exit_code == 0
     assert "Database migration successful!" in captured.out
 
-    system_commands.redis_client.lock.assert_called_once_with(name="db_upgrade_lock", timeout=999, thread_local=False)
+    _redis_client_mock().lock.assert_called_once_with(name="db_upgrade_lock", timeout=999, thread_local=False)
     lock.acquire.assert_called_once_with(blocking=False)
     lock.release.assert_called_once()
 
 
-def test_upgrade_db_renews_lock_during_migration(monkeypatch, capsys):
+def test_upgrade_db_renews_lock_during_migration(
+    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
+) -> None:
     """
     Ensure the lock is renewed while migrations are running, so the base TTL can stay short.
     """
@@ -405,17 +423,17 @@ def test_upgrade_db_renews_lock_during_migration(monkeypatch, capsys):
 
     lock = MagicMock()
     lock.acquire.return_value = True
-    system_commands.redis_client.lock.return_value = lock
+    _redis_client_mock().lock.return_value = lock
 
     renewed = threading.Event()
 
-    def _reacquire():
+    def _reacquire() -> bool:
         renewed.set()
         return True
 
     lock.reacquire.side_effect = _reacquire
 
-    def _upgrade():
+    def _upgrade() -> None:
         assert renewed.wait(HEARTBEAT_WAIT_TIMEOUT_SECONDS)
 
     _install_fake_flask_migrate(monkeypatch, _upgrade)
@@ -427,23 +445,25 @@ def test_upgrade_db_renews_lock_during_migration(monkeypatch, capsys):
     assert lock.reacquire.call_count >= 1
 
 
-def test_upgrade_db_ignores_reacquire_errors(monkeypatch, capsys):
+def test_upgrade_db_ignores_reacquire_errors(
+    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
+) -> None:
     # Use a small TTL so heartbeat runs during the upgrade call.
     monkeypatch.setattr(system_commands, "DB_UPGRADE_LOCK_TTL_SECONDS", 0.3)
 
     lock = MagicMock()
     lock.acquire.return_value = True
-    system_commands.redis_client.lock.return_value = lock
+    _redis_client_mock().lock.return_value = lock
 
     attempted = threading.Event()
 
-    def _reacquire():
+    def _reacquire() -> None:
         attempted.set()
         raise RedisError("simulated")
 
     lock.reacquire.side_effect = _reacquire
 
-    def _upgrade():
+    def _upgrade() -> None:
         assert attempted.wait(HEARTBEAT_WAIT_TIMEOUT_SECONDS)
 
     _install_fake_flask_migrate(monkeypatch, _upgrade)

api/tests/unit_tests/pyrefly.toml

@@ -8,9 +8,7 @@ min-severity = "warn"
 
 # Existing strict-mode debt. Remove a file when bringing it under strict checking.
 project-excludes = [
-  "commands/test_retention.py",
   "dev/test_generate_swagger_specs.py",
-  "commands/test_system.py",
   "controllers/console/app/test_agent_config_inspector.py",
   "controllers/console/app/test_app.py",
   "controllers/console/app/test_app_import.py",