From a50b6243ec118ffbb64cfff326d2e5f56d8e1e9a Mon Sep 17 00:00:00 2001 From: jtaroreh Date: Fri, 23 Oct 2026 10:41:05 -0700 Subject: [PATCH] Refactor(tests): Prune test bloat, consolidate provisioner suite, and fix steering collision retry - Prune uncollected dead code, duplicate class aliases, and non-hermetic checks across test suites - Relocate TestSetupRuntimeProvisioner into tests/test_setup_runtime.py - Fix concurrency collision handling in StorageMessenger.send_steer with outer retry loop and atomic O_CREAT | O_EXCL file creation - Restore tests 16, 17, and 18 inside TestConversationalSteeringTool - Prune orphaned imports and attributes in test_cloud_preflight_and_worker.py and test_container_spec.py - All 314 tests verified passing with zero regressions --- skills/swarm/scripts/storage_messenger.py | 123 ++++++---- tests/test_cloud_preflight_and_worker.py | 262 --------------------- tests/test_container_spec.py | 26 -- tests/test_conversational_steering_tool.py | 11 +- tests/test_reactive_coordination_rules.py | 124 ---------- tests/test_setup_runtime.py | 260 ++++++++++++++++++++ tests/test_storage_uploader.py | 8 - tests/test_swarm_generic.py | 3 - 8 files changed, 337 insertions(+), 480 deletions(-) diff --git a/skills/swarm/scripts/storage_messenger.py b/skills/swarm/scripts/storage_messenger.py index 668b01c..d4169e5 100644 --- a/skills/swarm/scripts/storage_messenger.py +++ b/skills/swarm/scripts/storage_messenger.py @@ -268,10 +268,17 @@ def upload_from_string( if_generation_match: Optional[int] = None, **kwargs: Any, ) -> None: - if if_generation_match == 0 and self.path.exists(): - raise FileExistsError(f"Target object already exists: {self.path}") self.path.parent.mkdir(parents=True, exist_ok=True) - self.path.write_text(data, encoding="utf-8") + if if_generation_match == 0: + flags = os.O_CREAT | os.O_EXCL | os.O_WRONLY + try: + fd = os.open(str(self.path), flags, 0o666) + except FileExistsError: + raise FileExistsError(f"Target object already exists: {self.path}") + with open(fd, "w", encoding="utf-8") as f: + f.write(data) + else: + self.path.write_text(data, encoding="utf-8") def download_as_text(self) -> str: return self.path.read_text(encoding="utf-8") @@ -563,62 +570,74 @@ def send_steer( retry_delay: float = 0.05, ) -> MailboxEnvelope: prefix = f"{self._session_prefix()}/tasks/task-{task_index}/inbox/" - max_seq = 0 + delay = retry_delay - def _resolve_max_seq(): - nonlocal max_seq - blobs = self.bucket.list_blobs(prefix=prefix) - for b in blobs: - name = getattr(b, "name", "") - if "_steer.json" in name: - base = name.split("/")[-1] - try: - seq_val = int(base.split("_")[0]) - if seq_val > max_seq: - max_seq = seq_val - except ValueError: - pass + for attempt in range(max_retries + 1): + max_seq = 0 - try: - retry_gcs_operation(_resolve_max_seq, max_retries=max_retries, initial_delay=retry_delay) - except Exception as exc: - raise RuntimeError("GCS sequence resolution failed") from exc + def _resolve_max_seq(): + nonlocal max_seq + blobs = self.bucket.list_blobs(prefix=prefix) + for b in blobs: + name = getattr(b, "name", "") + if "_steer.json" in name: + base = name.split("/")[-1] + try: + seq_val = int(base.split("_")[0]) + if seq_val > max_seq: + max_seq = seq_val + except ValueError: + pass - steer_seq = max_seq + 1 - envelope = MailboxEnvelope( - session_id=self.session_id, - task_index=task_index, - seq=steer_seq, - msg_type=MessageType.STEER, - sender="orchestrator", - recipient=f"worker-{task_index}", - timestamp=time.time(), - payload={"instruction": instruction, "text": instruction}, - status=MessageStatus.PENDING, - ) + try: + retry_gcs_operation(_resolve_max_seq, max_retries=max_retries, initial_delay=retry_delay) + except Exception as exc: + raise RuntimeError("GCS sequence resolution failed") from exc + + steer_seq = max_seq + 1 + envelope = MailboxEnvelope( + session_id=self.session_id, + task_index=task_index, + seq=steer_seq, + msg_type=MessageType.STEER, + sender="orchestrator", + recipient=f"worker-{task_index}", + timestamp=time.time(), + payload={"instruction": instruction, "text": instruction}, + status=MessageStatus.PENDING, + ) + + obj_path = self._inbox_steer_path(task_index, steer_seq) + blob = self.bucket.blob(obj_path) - obj_path = self._inbox_steer_path(task_index, steer_seq) - blob = self.bucket.blob(obj_path) + def _upload_blob(): + try: + blob.upload_from_string( + envelope.to_json(), + content_type="application/json", + if_generation_match=0, + ) + except TypeError: + blob.upload_from_string( + envelope.to_json(), + content_type="application/json", + ) - def _upload_blob(): try: - blob.upload_from_string( - envelope.to_json(), - content_type="application/json", - if_generation_match=0, + retry_gcs_operation( + _upload_blob, + max_retries=max_retries, + initial_delay=retry_delay, ) - except TypeError: - blob.upload_from_string( - envelope.to_json(), - content_type="application/json", - ) - - retry_gcs_operation( - _upload_blob, - max_retries=max_retries, - initial_delay=retry_delay, - ) - return envelope + return envelope + except Exception as exc: + if _is_precondition_or_collision_error(exc): + if attempt >= max_retries: + raise + time.sleep(delay) + delay *= 1.5 + continue + raise def discover_tasks(self) -> List[int]: """Discovers active or existing task indices for this session from storage.""" diff --git a/tests/test_cloud_preflight_and_worker.py b/tests/test_cloud_preflight_and_worker.py index 971ffaa..93d5df9 100644 --- a/tests/test_cloud_preflight_and_worker.py +++ b/tests/test_cloud_preflight_and_worker.py @@ -11,9 +11,7 @@ from unittest.mock import AsyncMock, MagicMock, patch sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "skills" / "swarm" / "scripts")) -sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "skills" / "setup-agystack" / "scripts")) import cloud_worker -import setup_runtime from cloud_dispatch import ( build_gcloud_command, format_monitored_line, @@ -1407,266 +1405,6 @@ def test_manifest_and_session_branch_resolution(self): self.assertEqual(branch_fallback, "worker-2") -class TestSetupRuntimeProvisioner(unittest.TestCase): - def setUp(self): - self.temp_dir_obj = tempfile.TemporaryDirectory() - self.temp_dir = Path(self.temp_dir_obj.name) - - def tearDown(self): - self.temp_dir_obj.cleanup() - - @patch("setup_runtime.run_cmd") - def test_enable_apis_includes_storage(self, mock_run_cmd): - mock_run_cmd.return_value = MagicMock(returncode=0, stdout="", stderr="") - setup_runtime.enable_apis("my-project-123") - mock_run_cmd.assert_called_once() - cmd = mock_run_cmd.call_args[0][0] - self.assertEqual(cmd[0:3], ["gcloud", "services", "enable"]) - self.assertIn("storage.googleapis.com", cmd) - self.assertIn("run.googleapis.com", cmd) - self.assertIn("artifactregistry.googleapis.com", cmd) - self.assertIn("cloudbuild.googleapis.com", cmd) - self.assertIn("aiplatform.googleapis.com", cmd) - self.assertIn("--project=my-project-123", cmd) - - @patch("setup_runtime.run_cmd") - def test_ensure_gcs_bucket_creates_if_not_exists(self, mock_run_cmd): - describe_res = MagicMock(returncode=1, stdout="", stderr="Bucket not found") - create_res = MagicMock(returncode=0, stdout="", stderr="") - mock_run_cmd.side_effect = [describe_res, create_res] - - bucket_name = setup_runtime.ensure_gcs_bucket("my-new-bucket", "my-project-123", "us-central1") - self.assertEqual(bucket_name, "my-new-bucket") - self.assertEqual(mock_run_cmd.call_count, 2) - describe_cmd = mock_run_cmd.call_args_list[0][0][0] - self.assertEqual(describe_cmd, ["gcloud", "storage", "buckets", "describe", "gs://my-new-bucket"]) - create_cmd = mock_run_cmd.call_args_list[1][0][0] - self.assertEqual( - create_cmd, - ["gcloud", "storage", "buckets", "create", "gs://my-new-bucket", "--project=my-project-123", "--location=us-central1"] - ) - - @patch("setup_runtime.run_cmd") - def test_ensure_gcs_bucket_skips_if_exists(self, mock_run_cmd): - describe_res = MagicMock(returncode=0, stdout="{}", stderr="") - mock_run_cmd.return_value = describe_res - - bucket_name = setup_runtime.ensure_gcs_bucket("gs://my-existing-bucket", "my-project-123", "us-central1") - self.assertEqual(bucket_name, "my-existing-bucket") - self.assertEqual(mock_run_cmd.call_count, 1) - describe_cmd = mock_run_cmd.call_args[0][0] - self.assertEqual(describe_cmd, ["gcloud", "storage", "buckets", "describe", "gs://my-existing-bucket"]) - - @patch("setup_runtime.run_cmd") - def test_configure_iam_permissions_default_compute_sa(self, mock_run_cmd): - project_desc_res = MagicMock(returncode=0, stdout="9876543210\n", stderr="") - bind_aiplatform_res = MagicMock(returncode=0, stdout="", stderr="") - bind_storage_res = MagicMock(returncode=0, stdout="", stderr="") - bind_secret_res = MagicMock(returncode=0, stdout="", stderr="") - mock_run_cmd.side_effect = [project_desc_res, bind_aiplatform_res, bind_storage_res, bind_secret_res] - - sa = setup_runtime.configure_iam_permissions("my-project-123") - self.assertEqual(sa, "9876543210-compute@developer.gserviceaccount.com") - self.assertEqual(mock_run_cmd.call_count, 4) - - desc_cmd = mock_run_cmd.call_args_list[0][0][0] - self.assertEqual(desc_cmd, ["gcloud", "projects", "describe", "my-project-123", "--format=value(projectNumber)"]) - - bind1 = mock_run_cmd.call_args_list[1][0][0] - self.assertEqual( - bind1, - [ - "gcloud", "projects", "add-iam-policy-binding", "my-project-123", - "--member=serviceAccount:9876543210-compute@developer.gserviceaccount.com", - "--role=roles/aiplatform.user" - ] - ) - - bind2 = mock_run_cmd.call_args_list[2][0][0] - self.assertEqual( - bind2, - [ - "gcloud", "projects", "add-iam-policy-binding", "my-project-123", - "--member=serviceAccount:9876543210-compute@developer.gserviceaccount.com", - "--role=roles/storage.objectAdmin" - ] - ) - - bind3 = mock_run_cmd.call_args_list[3][0][0] - self.assertEqual( - bind3, - [ - "gcloud", "projects", "add-iam-policy-binding", "my-project-123", - "--member=serviceAccount:9876543210-compute@developer.gserviceaccount.com", - "--role=roles/secretmanager.secretAccessor" - ] - ) - - @patch("setup_runtime.run_cmd") - def test_configure_iam_permissions_custom_sa(self, mock_run_cmd): - bind_aiplatform_res = MagicMock(returncode=0, stdout="", stderr="") - bind_storage_res = MagicMock(returncode=0, stdout="", stderr="") - bind_secret_res = MagicMock(returncode=0, stdout="", stderr="") - mock_run_cmd.side_effect = [bind_aiplatform_res, bind_storage_res, bind_secret_res] - - sa = setup_runtime.configure_iam_permissions("my-project-123", "custom-runner@my-project-123.iam.gserviceaccount.com") - self.assertEqual(sa, "custom-runner@my-project-123.iam.gserviceaccount.com") - self.assertEqual(mock_run_cmd.call_count, 3) - - bind1 = mock_run_cmd.call_args_list[0][0][0] - self.assertEqual( - bind1, - [ - "gcloud", "projects", "add-iam-policy-binding", "my-project-123", - "--member=serviceAccount:custom-runner@my-project-123.iam.gserviceaccount.com", - "--role=roles/aiplatform.user" - ] - ) - - bind2 = mock_run_cmd.call_args_list[1][0][0] - self.assertEqual( - bind2, - [ - "gcloud", "projects", "add-iam-policy-binding", "my-project-123", - "--member=serviceAccount:custom-runner@my-project-123.iam.gserviceaccount.com", - "--role=roles/storage.objectAdmin" - ] - ) - - bind3 = mock_run_cmd.call_args_list[2][0][0] - self.assertEqual( - bind3, - [ - "gcloud", "projects", "add-iam-policy-binding", "my-project-123", - "--member=serviceAccount:custom-runner@my-project-123.iam.gserviceaccount.com", - "--role=roles/secretmanager.secretAccessor" - ] - ) - - @patch("setup_runtime.run_cmd") - def test_build_and_deploy_worker_creates_job_with_memory_and_cpu(self, mock_run_cmd): - repo_desc_res = MagicMock(returncode=0, stdout="", stderr="") - build_res = MagicMock(returncode=0, stdout="", stderr="") - job_desc_res = MagicMock(returncode=1, stdout="", stderr="Job not found") - job_create_res = MagicMock(returncode=0, stdout="", stderr="") - mock_run_cmd.side_effect = [repo_desc_res, build_res, job_desc_res, job_create_res] - - setup_runtime.build_and_deploy_worker( - project_id="my-project-123", - region="us-central1", - image_tag="us-central1-docker.pkg.dev/my-project-123/agystack/worker:latest", - scripts_dir="skills/swarm/scripts", - service_account="custom-sa@example.com", - ) - - job_cmd = mock_run_cmd.call_args_list[-1][0][0] - self.assertEqual(job_cmd[0:4], ["gcloud", "run", "jobs", "create"]) - self.assertIn("--memory=2Gi", job_cmd) - self.assertIn("--cpu=2", job_cmd) - self.assertIn("--parallelism=16", job_cmd) - self.assertIn("--max-retries=0", job_cmd) - self.assertIn("--image=us-central1-docker.pkg.dev/my-project-123/agystack/worker:latest", job_cmd) - self.assertIn("--region=us-central1", job_cmd) - self.assertIn("--project=my-project-123", job_cmd) - self.assertIn("--service-account=custom-sa@example.com", job_cmd) - - @patch("setup_runtime.run_cmd") - def test_build_and_deploy_worker_updates_job_with_memory_and_cpu(self, mock_run_cmd): - repo_desc_res = MagicMock(returncode=0, stdout="", stderr="") - build_res = MagicMock(returncode=0, stdout="", stderr="") - job_desc_res = MagicMock(returncode=0, stdout="{}", stderr="") - job_update_res = MagicMock(returncode=0, stdout="", stderr="") - mock_run_cmd.side_effect = [repo_desc_res, build_res, job_desc_res, job_update_res] - - setup_runtime.build_and_deploy_worker( - project_id="my-project-123", - region="us-central1", - image_tag="us-central1-docker.pkg.dev/my-project-123/agystack/worker:latest", - scripts_dir="skills/swarm/scripts", - service_account="custom-sa@example.com", - ) - - job_cmd = mock_run_cmd.call_args_list[-1][0][0] - self.assertEqual(job_cmd[0:4], ["gcloud", "run", "jobs", "update"]) - self.assertIn("--memory=2Gi", job_cmd) - self.assertIn("--cpu=2", job_cmd) - self.assertIn("--parallelism=16", job_cmd) - self.assertIn("--max-retries=0", job_cmd) - self.assertIn("--service-account=custom-sa@example.com", job_cmd) - - def test_write_runtime_config_persists_gcs_bucket(self): - target_file = self.temp_dir / "agystack-runtime.json" - config = setup_runtime.write_runtime_config( - project_id="test-proj", - region="us-central1", - job_name="agystack-swarm-worker", - image_tag="us-central1-docker.pkg.dev/test-proj/agystack/worker:latest", - auth_mode="vertex", - gcs_bucket="test-proj-swarm-results", - target_paths=[target_file], - service_account="test-sa@developer.gserviceaccount.com", - ) - - self.assertEqual(config["gcs_bucket"], "test-proj-swarm-results") - self.assertEqual(config["auth_mode"], "vertex") - self.assertEqual(config["runtime"], "cloud-run") - self.assertEqual(config["service_account"], "test-sa@developer.gserviceaccount.com") - - loaded = json.loads(target_file.read_text(encoding="utf-8")) - self.assertEqual(loaded["gcs_bucket"], "test-proj-swarm-results") - self.assertEqual(loaded["job_name"], "agystack-swarm-worker") - self.assertEqual(loaded["parallelism"], 16) - self.assertEqual(loaded["service_account"], "test-sa@developer.gserviceaccount.com") - - @patch("setup_runtime.write_runtime_config") - @patch("setup_runtime.run_cmd") - def test_run_provisioning_full_flow(self, mock_run_cmd, mock_write_cfg): - mock_run_cmd.side_effect = [ - MagicMock(returncode=0, stdout="", stderr=""), # enable_apis - MagicMock(returncode=1, stdout="", stderr=""), # bucket describe - MagicMock(returncode=0, stdout="", stderr=""), # bucket create - MagicMock(returncode=0, stdout="", stderr=""), # role binding aiplatform - MagicMock(returncode=0, stdout="", stderr=""), # role binding storage - MagicMock(returncode=0, stdout="", stderr=""), # role binding secretmanager - MagicMock(returncode=0, stdout="", stderr=""), # repo describe - MagicMock(returncode=0, stdout="", stderr=""), # builds submit - MagicMock(returncode=1, stdout="", stderr=""), # job describe (1 => create) - MagicMock(returncode=0, stdout="", stderr=""), # job create - ] - - summary = setup_runtime.run_provisioning( - project_id="test-proj-full", - region="us-central1", - scripts_dir="skills/swarm/scripts", - bucket_name="custom-bucket", - service_account="custom-sa@developer.gserviceaccount.com", - ) - - all_cmds = [call[0][0] for call in mock_run_cmd.call_args_list] - - api_cmd = next(cmd for cmd in all_cmds if cmd[0:3] == ["gcloud", "services", "enable"]) - self.assertIn("storage.googleapis.com", api_cmd) - - bucket_create_cmd = next(cmd for cmd in all_cmds if len(cmd) > 3 and cmd[0:4] == ["gcloud", "storage", "buckets", "create"]) - self.assertIn("gs://custom-bucket", bucket_create_cmd) - self.assertIn("--project=test-proj-full", bucket_create_cmd) - - iam_cmds = [cmd for cmd in all_cmds if len(cmd) > 3 and cmd[0:3] == ["gcloud", "projects", "add-iam-policy-binding"]] - self.assertEqual(len(iam_cmds), 3) - roles_bound = {next(arg.split("=")[1] for arg in cmd if arg.startswith("--role=")) for cmd in iam_cmds} - self.assertEqual(roles_bound, {"roles/aiplatform.user", "roles/storage.objectAdmin", "roles/secretmanager.secretAccessor"}) - - job_create_cmd = next(cmd for cmd in all_cmds if len(cmd) > 3 and cmd[0:4] == ["gcloud", "run", "jobs", "create"]) - self.assertIn("--memory=2Gi", job_create_cmd) - self.assertIn("--cpu=2", job_create_cmd) - self.assertIn("--service-account=custom-sa@developer.gserviceaccount.com", job_create_cmd) - - mock_write_cfg.assert_called_once() - self.assertEqual(mock_write_cfg.call_args[1].get("gcs_bucket"), "custom-bucket") - self.assertEqual(mock_write_cfg.call_args[1].get("service_account"), "custom-sa@developer.gserviceaccount.com") - self.assertEqual(summary.get("service_account"), "custom-sa@developer.gserviceaccount.com") - - class TestDispatcherPartialFailureGating(unittest.TestCase): def test_partial_failure_output_gating_warning_and_suppression(self): from cloud_dispatch import print_candidate_patches_table diff --git a/tests/test_container_spec.py b/tests/test_container_spec.py index a61a82c..0c672ac 100644 --- a/tests/test_container_spec.py +++ b/tests/test_container_spec.py @@ -1,4 +1,3 @@ -import os import re import unittest from pathlib import Path @@ -10,7 +9,6 @@ def setUpClass(cls): cls.project_root = Path(__file__).resolve().parent.parent cls.scripts_dir = cls.project_root / "skills" / "swarm" / "scripts" cls.dockerfile_path = cls.scripts_dir / "Dockerfile" - cls.installed_plugin_dir = Path(os.path.expanduser("~/.gemini/config/plugins/agystack/skills/swarm/scripts")) if not cls.dockerfile_path.is_file(): raise FileNotFoundError(f"Dockerfile not found at {cls.dockerfile_path}") @@ -120,20 +118,6 @@ def test_plugin_skills_and_rules_copied(self): r"COPY\s+agents\s+\/home\/worker\/\.gemini\/config\/plugins\/agystack\/agents", ) - def test_cargo_deny_does_not_exist_on_disk(self): - cargo_deny_path = self.scripts_dir / "cargo-deny" - self.assertFalse( - cargo_deny_path.exists(), - f"cargo-deny must not exist on disk: {cargo_deny_path}", - ) - - def test_setup_rust_worker_does_not_exist_on_disk(self): - setup_script_path = self.scripts_dir / "setup-rust-worker.sh" - self.assertFalse( - setup_script_path.exists(), - f"setup-rust-worker.sh must not exist on disk: {setup_script_path}", - ) - def test_cargo_deny_not_in_dockerfile(self): self.assertNotIn( "cargo-deny", @@ -141,16 +125,6 @@ def test_cargo_deny_not_in_dockerfile(self): "cargo-deny must not be referenced in Dockerfile", ) - def test_installed_plugin_dockerfile_sync(self): - installed_dockerfile = self.installed_plugin_dir / "Dockerfile" - if installed_dockerfile.is_file(): - installed_content = installed_dockerfile.read_text(encoding="utf-8") - self.assertEqual( - self.dockerfile_content, - installed_content, - "Installed agystack plugin Dockerfile must match project Dockerfile", - ) - if __name__ == "__main__": unittest.main() diff --git a/tests/test_conversational_steering_tool.py b/tests/test_conversational_steering_tool.py index 3e46d4a..aacc777 100644 --- a/tests/test_conversational_steering_tool.py +++ b/tests/test_conversational_steering_tool.py @@ -549,10 +549,6 @@ def test_15_steering_context_preserved_alongside_stagnancy_warning(self): self.assertIn(steer_text, res.custom_context) self.assertIn("consecutive times with identical arguments", res.custom_context) - -if __name__ == "__main__": - unittest.main() - def test_16_concurrent_send_steer_collision_retry(self): """Simulate two simultaneous send_steer() calls colliding on sequence number, verifying both succeed with distinct sequence numbers.""" @@ -613,7 +609,7 @@ def mock_upload(*args, **kwargs): nonlocal call_count call_count += 1 if call_count == 1: - orch.bucket.blob(name).upload_from_string("{}", content_type="application/json") + real_blob_factory(name).upload_from_string("{}", content_type="application/json") raise PreconditionFailed("Simulated HTTP 412 Precondition Failed") return orig_upload(*args, **kwargs) @@ -637,3 +633,8 @@ def test_18_mailbox_controller_modular_import(self): self.assertTrue(callable(mc_list)) self.assertTrue(callable(mc_reply)) self.assertTrue(callable(mc_steer)) + + +if __name__ == "__main__": + unittest.main() + diff --git a/tests/test_reactive_coordination_rules.py b/tests/test_reactive_coordination_rules.py index 6b35956..b4e57ee 100644 --- a/tests/test_reactive_coordination_rules.py +++ b/tests/test_reactive_coordination_rules.py @@ -8,80 +8,6 @@ SKILLS_DIR = REPO_ROOT / "skills" -class CoordinatorStateMachine: - """Models coordinator state transitions during subagent fan-out and reactive wakeup.""" - - def __init__(self, expected_workers: int, timeout_seconds: int = 300): - self.expected_workers = expected_workers - self.timeout_seconds = timeout_seconds - self.state = "INIT" - self.watchdog_armed = False - self.watchdog_timer_id = None - self.results = {} - self.stragglers = [] - self.terminated_stragglers = [] - - def dispatch(self, timer_task_id: str = "task-watchdog-1"): - """Turn 1: Dispatch subagents and arm unconditional global watchdog timer.""" - self.state = "DISPATCHED" - self.watchdog_armed = True - self.watchdog_timer_id = timer_task_id - # Turn 2: yield turn with zero tool calls - self.state = "WAITING" - yield_tool_calls = [] # ZERO tool calls - return yield_tool_calls - - def on_worker_message(self, worker_id: str, findings: str): - """Worker completion reactively wakes the coordinator.""" - assert self.state in ("WAITING", "PARTIAL_RECEIVED") - self.results[worker_id] = findings - - if len(self.results) < self.expected_workers: - self.state = "PARTIAL_RECEIVED" - # Watchdog remains armed! Yield immediately with zero tool calls - return {"action": "YIELD", "tool_calls": []} - else: - self.state = "COMPLETED" - # Cancel watchdog timer and proceed to synthesis - timer_to_kill = self.watchdog_timer_id - self.watchdog_armed = False - self.watchdog_timer_id = None - return { - "action": "CANCEL_WATCHDOG_AND_SYNTHESIZE", - "tool_calls": [ - { - "name": "manage_task", - "args": {"Action": "kill", "TaskId": timer_to_kill}, - } - ], - } - - def on_watchdog_timeout(self, active_subagents_list: list): - """Watchdog fires on timeout.""" - assert self.state in ("WAITING", "PARTIAL_RECEIVED") - self.state = "TIMED_OUT" - self.watchdog_armed = False - - # Timeout diagnosis: inspect manage_subagents(list) - missing = [w for w in active_subagents_list if w["id"] not in self.results] - self.stragglers = missing - self.terminated_stragglers = [w["id"] for w in missing] - - # Kill stragglers and proceed with partial results - kill_tool_calls = [ - { - "name": "manage_subagents", - "args": {"Action": "kill", "ConversationIds": [w["id"]]}, - } - for w in missing - ] - return { - "action": "SYNTHESIZE_PARTIAL", - "tool_calls": kill_tool_calls, - "results": self.results, - } - - class TestReactiveCoordinationRules(unittest.TestCase): def test_agents_md_has_reactive_wakeup_invariant(self): agents_file = RULES_DIR / "AGENTS.md" @@ -143,56 +69,6 @@ def test_arena_and_swarm_skills_have_reactive_protocol(self): self.assertIn("send_message", swarm_content) self.assertIn("{PARENT_CONVERSATION_ID}", swarm_content) - def test_coordinator_state_machine_logic(self): - sm = CoordinatorStateMachine(expected_workers=3, timeout_seconds=300) - - tool_calls = sm.dispatch("task-wd-1") - self.assertEqual(sm.state, "WAITING") - self.assertTrue(sm.watchdog_armed) - self.assertEqual(tool_calls, []) - - res1 = sm.on_worker_message("worker-1", "Finding 1: bug in parser") - self.assertEqual(sm.state, "PARTIAL_RECEIVED") - self.assertTrue(sm.watchdog_armed) - self.assertEqual(res1["action"], "YIELD") - self.assertEqual(res1["tool_calls"], []) - - res2 = sm.on_worker_message("worker-2", "Finding 2: race condition") - self.assertEqual(sm.state, "PARTIAL_RECEIVED") - self.assertTrue(sm.watchdog_armed) - self.assertEqual(res2["action"], "YIELD") - self.assertEqual(res2["tool_calls"], []) - - res3 = sm.on_worker_message("worker-3", "No findings.") - self.assertEqual(sm.state, "COMPLETED") - self.assertFalse(sm.watchdog_armed) - self.assertEqual(res3["action"], "CANCEL_WATCHDOG_AND_SYNTHESIZE") - self.assertEqual(len(res3["tool_calls"]), 1) - self.assertEqual(res3["tool_calls"][0]["name"], "manage_task") - self.assertEqual(res3["tool_calls"][0]["args"]["Action"], "kill") - self.assertEqual(res3["tool_calls"][0]["args"]["TaskId"], "task-wd-1") - - sm2 = CoordinatorStateMachine(expected_workers=3, timeout_seconds=300) - sm2.dispatch("task-wd-2") - sm2.on_worker_message("worker-1", "Finding A") - sm2.on_worker_message("worker-2", "Finding B") - - active_workers = [ - {"id": "worker-1", "status": "done"}, - {"id": "worker-2", "status": "done"}, - {"id": "worker-3", "status": "idle"}, - ] - timeout_res = sm2.on_watchdog_timeout(active_workers) - self.assertEqual(sm2.state, "TIMED_OUT") - self.assertEqual(timeout_res["action"], "SYNTHESIZE_PARTIAL") - self.assertEqual(sm2.terminated_stragglers, ["worker-3"]) - self.assertEqual(len(timeout_res["tool_calls"]), 1) - self.assertEqual(timeout_res["tool_calls"][0]["name"], "manage_subagents") - self.assertEqual( - timeout_res["tool_calls"][0]["args"]["ConversationIds"], ["worker-3"] - ) - self.assertEqual(len(timeout_res["results"]), 2) - def test_agents_md_has_delegation_thresholds(self): agents_file = RULES_DIR / "AGENTS.md" content = agents_file.read_text(encoding="utf-8") diff --git a/tests/test_setup_runtime.py b/tests/test_setup_runtime.py index 997c86f..636c7b4 100644 --- a/tests/test_setup_runtime.py +++ b/tests/test_setup_runtime.py @@ -569,5 +569,265 @@ def test_doctor_platform_failure_exits_1(self): self.assertFalse(report["platform"]["ok"]) +class TestSetupRuntimeProvisioner(unittest.TestCase): + def setUp(self): + self.temp_dir_obj = tempfile.TemporaryDirectory() + self.temp_dir = Path(self.temp_dir_obj.name) + + def tearDown(self): + self.temp_dir_obj.cleanup() + + @patch("setup_runtime.run_cmd") + def test_enable_apis_includes_storage(self, mock_run_cmd): + mock_run_cmd.return_value = MagicMock(returncode=0, stdout="", stderr="") + setup_runtime.enable_apis("my-project-123") + mock_run_cmd.assert_called_once() + cmd = mock_run_cmd.call_args[0][0] + self.assertEqual(cmd[0:3], ["gcloud", "services", "enable"]) + self.assertIn("storage.googleapis.com", cmd) + self.assertIn("run.googleapis.com", cmd) + self.assertIn("artifactregistry.googleapis.com", cmd) + self.assertIn("cloudbuild.googleapis.com", cmd) + self.assertIn("aiplatform.googleapis.com", cmd) + self.assertIn("--project=my-project-123", cmd) + + @patch("setup_runtime.run_cmd") + def test_ensure_gcs_bucket_creates_if_not_exists(self, mock_run_cmd): + describe_res = MagicMock(returncode=1, stdout="", stderr="Bucket not found") + create_res = MagicMock(returncode=0, stdout="", stderr="") + mock_run_cmd.side_effect = [describe_res, create_res] + + bucket_name = setup_runtime.ensure_gcs_bucket("my-new-bucket", "my-project-123", "us-central1") + self.assertEqual(bucket_name, "my-new-bucket") + self.assertEqual(mock_run_cmd.call_count, 2) + describe_cmd = mock_run_cmd.call_args_list[0][0][0] + self.assertEqual(describe_cmd, ["gcloud", "storage", "buckets", "describe", "gs://my-new-bucket"]) + create_cmd = mock_run_cmd.call_args_list[1][0][0] + self.assertEqual( + create_cmd, + ["gcloud", "storage", "buckets", "create", "gs://my-new-bucket", "--project=my-project-123", "--location=us-central1"] + ) + + @patch("setup_runtime.run_cmd") + def test_ensure_gcs_bucket_skips_if_exists(self, mock_run_cmd): + describe_res = MagicMock(returncode=0, stdout="{}", stderr="") + mock_run_cmd.return_value = describe_res + + bucket_name = setup_runtime.ensure_gcs_bucket("gs://my-existing-bucket", "my-project-123", "us-central1") + self.assertEqual(bucket_name, "my-existing-bucket") + self.assertEqual(mock_run_cmd.call_count, 1) + describe_cmd = mock_run_cmd.call_args[0][0] + self.assertEqual(describe_cmd, ["gcloud", "storage", "buckets", "describe", "gs://my-existing-bucket"]) + + @patch("setup_runtime.run_cmd") + def test_configure_iam_permissions_default_compute_sa(self, mock_run_cmd): + project_desc_res = MagicMock(returncode=0, stdout="9876543210\n", stderr="") + bind_aiplatform_res = MagicMock(returncode=0, stdout="", stderr="") + bind_storage_res = MagicMock(returncode=0, stdout="", stderr="") + bind_secret_res = MagicMock(returncode=0, stdout="", stderr="") + mock_run_cmd.side_effect = [project_desc_res, bind_aiplatform_res, bind_storage_res, bind_secret_res] + + sa = setup_runtime.configure_iam_permissions("my-project-123") + self.assertEqual(sa, "9876543210-compute@developer.gserviceaccount.com") + self.assertEqual(mock_run_cmd.call_count, 4) + + desc_cmd = mock_run_cmd.call_args_list[0][0][0] + self.assertEqual(desc_cmd, ["gcloud", "projects", "describe", "my-project-123", "--format=value(projectNumber)"]) + + bind1 = mock_run_cmd.call_args_list[1][0][0] + self.assertEqual( + bind1, + [ + "gcloud", "projects", "add-iam-policy-binding", "my-project-123", + "--member=serviceAccount:9876543210-compute@developer.gserviceaccount.com", + "--role=roles/aiplatform.user" + ] + ) + + bind2 = mock_run_cmd.call_args_list[2][0][0] + self.assertEqual( + bind2, + [ + "gcloud", "projects", "add-iam-policy-binding", "my-project-123", + "--member=serviceAccount:9876543210-compute@developer.gserviceaccount.com", + "--role=roles/storage.objectAdmin" + ] + ) + + bind3 = mock_run_cmd.call_args_list[3][0][0] + self.assertEqual( + bind3, + [ + "gcloud", "projects", "add-iam-policy-binding", "my-project-123", + "--member=serviceAccount:9876543210-compute@developer.gserviceaccount.com", + "--role=roles/secretmanager.secretAccessor" + ] + ) + + @patch("setup_runtime.run_cmd") + def test_configure_iam_permissions_custom_sa(self, mock_run_cmd): + bind_aiplatform_res = MagicMock(returncode=0, stdout="", stderr="") + bind_storage_res = MagicMock(returncode=0, stdout="", stderr="") + bind_secret_res = MagicMock(returncode=0, stdout="", stderr="") + mock_run_cmd.side_effect = [bind_aiplatform_res, bind_storage_res, bind_secret_res] + + sa = setup_runtime.configure_iam_permissions("my-project-123", "custom-runner@my-project-123.iam.gserviceaccount.com") + self.assertEqual(sa, "custom-runner@my-project-123.iam.gserviceaccount.com") + self.assertEqual(mock_run_cmd.call_count, 3) + + bind1 = mock_run_cmd.call_args_list[0][0][0] + self.assertEqual( + bind1, + [ + "gcloud", "projects", "add-iam-policy-binding", "my-project-123", + "--member=serviceAccount:custom-runner@my-project-123.iam.gserviceaccount.com", + "--role=roles/aiplatform.user" + ] + ) + + bind2 = mock_run_cmd.call_args_list[1][0][0] + self.assertEqual( + bind2, + [ + "gcloud", "projects", "add-iam-policy-binding", "my-project-123", + "--member=serviceAccount:custom-runner@my-project-123.iam.gserviceaccount.com", + "--role=roles/storage.objectAdmin" + ] + ) + + bind3 = mock_run_cmd.call_args_list[2][0][0] + self.assertEqual( + bind3, + [ + "gcloud", "projects", "add-iam-policy-binding", "my-project-123", + "--member=serviceAccount:custom-runner@my-project-123.iam.gserviceaccount.com", + "--role=roles/secretmanager.secretAccessor" + ] + ) + + @patch("setup_runtime.run_cmd") + def test_build_and_deploy_worker_creates_job_with_memory_and_cpu(self, mock_run_cmd): + repo_desc_res = MagicMock(returncode=0, stdout="", stderr="") + build_res = MagicMock(returncode=0, stdout="", stderr="") + job_desc_res = MagicMock(returncode=1, stdout="", stderr="Job not found") + job_create_res = MagicMock(returncode=0, stdout="", stderr="") + mock_run_cmd.side_effect = [repo_desc_res, build_res, job_desc_res, job_create_res] + + setup_runtime.build_and_deploy_worker( + project_id="my-project-123", + region="us-central1", + image_tag="us-central1-docker.pkg.dev/my-project-123/agystack/worker:latest", + scripts_dir="skills/swarm/scripts", + service_account="custom-sa@example.com", + ) + + job_cmd = mock_run_cmd.call_args_list[-1][0][0] + self.assertEqual(job_cmd[0:4], ["gcloud", "run", "jobs", "create"]) + self.assertIn("--memory=2Gi", job_cmd) + self.assertIn("--cpu=2", job_cmd) + self.assertIn("--parallelism=16", job_cmd) + self.assertIn("--max-retries=0", job_cmd) + self.assertIn("--image=us-central1-docker.pkg.dev/my-project-123/agystack/worker:latest", job_cmd) + self.assertIn("--region=us-central1", job_cmd) + self.assertIn("--project=my-project-123", job_cmd) + self.assertIn("--service-account=custom-sa@example.com", job_cmd) + + @patch("setup_runtime.run_cmd") + def test_build_and_deploy_worker_updates_job_with_memory_and_cpu(self, mock_run_cmd): + repo_desc_res = MagicMock(returncode=0, stdout="", stderr="") + build_res = MagicMock(returncode=0, stdout="", stderr="") + job_desc_res = MagicMock(returncode=0, stdout="{}", stderr="") + job_update_res = MagicMock(returncode=0, stdout="", stderr="") + mock_run_cmd.side_effect = [repo_desc_res, build_res, job_desc_res, job_update_res] + + setup_runtime.build_and_deploy_worker( + project_id="my-project-123", + region="us-central1", + image_tag="us-central1-docker.pkg.dev/my-project-123/agystack/worker:latest", + scripts_dir="skills/swarm/scripts", + service_account="custom-sa@example.com", + ) + + job_cmd = mock_run_cmd.call_args_list[-1][0][0] + self.assertEqual(job_cmd[0:4], ["gcloud", "run", "jobs", "update"]) + self.assertIn("--memory=2Gi", job_cmd) + self.assertIn("--cpu=2", job_cmd) + self.assertIn("--parallelism=16", job_cmd) + self.assertIn("--max-retries=0", job_cmd) + self.assertIn("--service-account=custom-sa@example.com", job_cmd) + + def test_write_runtime_config_persists_gcs_bucket(self): + target_file = self.temp_dir / "agystack-runtime.json" + config = setup_runtime.write_runtime_config( + project_id="test-proj", + region="us-central1", + job_name="agystack-swarm-worker", + image_tag="us-central1-docker.pkg.dev/test-proj/agystack/worker:latest", + auth_mode="vertex", + gcs_bucket="test-proj-swarm-results", + target_paths=[target_file], + service_account="test-sa@developer.gserviceaccount.com", + ) + + self.assertEqual(config["gcs_bucket"], "test-proj-swarm-results") + self.assertEqual(config["auth_mode"], "vertex") + self.assertEqual(config["runtime"], "cloud-run") + self.assertEqual(config["service_account"], "test-sa@developer.gserviceaccount.com") + + loaded = json.loads(target_file.read_text(encoding="utf-8")) + self.assertEqual(loaded["gcs_bucket"], "test-proj-swarm-results") + self.assertEqual(loaded["job_name"], "agystack-swarm-worker") + self.assertEqual(loaded["parallelism"], 16) + self.assertEqual(loaded["service_account"], "test-sa@developer.gserviceaccount.com") + + @patch("setup_runtime.write_runtime_config") + @patch("setup_runtime.run_cmd") + def test_run_provisioning_full_flow(self, mock_run_cmd, mock_write_cfg): + mock_run_cmd.side_effect = [ + MagicMock(returncode=0, stdout="", stderr=""), # enable_apis + MagicMock(returncode=1, stdout="", stderr=""), # bucket describe + MagicMock(returncode=0, stdout="", stderr=""), # bucket create + MagicMock(returncode=0, stdout="", stderr=""), # role binding aiplatform + MagicMock(returncode=0, stdout="", stderr=""), # role binding storage + MagicMock(returncode=0, stdout="", stderr=""), # role binding secretmanager + MagicMock(returncode=0, stdout="", stderr=""), # repo describe + MagicMock(returncode=0, stdout="", stderr=""), # builds submit + MagicMock(returncode=1, stdout="", stderr=""), # job describe (1 => create) + MagicMock(returncode=0, stdout="", stderr=""), # job create + ] + + summary = setup_runtime.run_provisioning( + project_id="test-proj-full", + region="us-central1", + scripts_dir="skills/swarm/scripts", + bucket_name="custom-bucket", + service_account="custom-sa@developer.gserviceaccount.com", + ) + + all_cmds = [call[0][0] for call in mock_run_cmd.call_args_list] + + api_cmd = next(cmd for cmd in all_cmds if cmd[0:3] == ["gcloud", "services", "enable"]) + self.assertIn("storage.googleapis.com", api_cmd) + + bucket_create_cmd = next(cmd for cmd in all_cmds if len(cmd) > 3 and cmd[0:4] == ["gcloud", "storage", "buckets", "create"]) + self.assertIn("gs://custom-bucket", bucket_create_cmd) + self.assertIn("--project=test-proj-full", bucket_create_cmd) + + iam_cmds = [cmd for cmd in all_cmds if len(cmd) > 3 and cmd[0:3] == ["gcloud", "projects", "add-iam-policy-binding"]] + self.assertEqual(len(iam_cmds), 3) + roles_bound = {next(arg.split("=")[1] for arg in cmd if arg.startswith("--role=")) for cmd in iam_cmds} + self.assertEqual(roles_bound, {"roles/aiplatform.user", "roles/storage.objectAdmin", "roles/secretmanager.secretAccessor"}) + + job_create_cmd = next(cmd for cmd in all_cmds if len(cmd) > 3 and cmd[0:4] == ["gcloud", "run", "jobs", "create"]) + self.assertIn("--memory=2Gi", job_create_cmd) + self.assertIn("--cpu=2", job_create_cmd) + self.assertIn("--service-account=custom-sa@developer.gserviceaccount.com", job_create_cmd) + + mock_write_cfg.assert_called_once() + self.assertEqual(mock_write_cfg.call_args[1].get("gcs_bucket"), "custom-bucket") + self.assertEqual(mock_write_cfg.call_args[1].get("service_account"), "custom-sa@developer.gserviceaccount.com") + self.assertEqual(summary.get("service_account"), "custom-sa@developer.gserviceaccount.com") + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_storage_uploader.py b/tests/test_storage_uploader.py index 45b5451..09ca080 100644 --- a/tests/test_storage_uploader.py +++ b/tests/test_storage_uploader.py @@ -27,11 +27,6 @@ def test_authenticated_url_shorthand(self): url = cloud_worker.build_authenticated_git_url("my-org/my-repo", "test_tok") self.assertEqual(url, "https://x-access-token:test_tok@github.com/my-org/my-repo.git") - def test_fork_repo_override_logic(self): - with patch.dict(os.environ, {"FORK_REPO_URL": "https://github.com/myfork/matrices-fast.git"}): - fork_url = os.environ.get("FORK_REPO_URL", "").strip() - self.assertEqual(fork_url, "https://github.com/myfork/matrices-fast.git") - class TestArtifactDiscovery(unittest.TestCase): def setUp(self): @@ -203,8 +198,5 @@ def test_parse_score_metrics_missing(self): self.assertIsNone(latency_ms) -TestStructuredOutputParsing = TestScoreParsing - - if __name__ == "__main__": unittest.main() diff --git a/tests/test_swarm_generic.py b/tests/test_swarm_generic.py index c742a7f..96edf3c 100644 --- a/tests/test_swarm_generic.py +++ b/tests/test_swarm_generic.py @@ -125,8 +125,5 @@ def test_empty_and_single_item_inputs(self): self.assertEqual(extract_pareto_frontier(single_min, maximize=False), single_min) -TestSwarmGenericHarvester = TestParetoFrontier - - if __name__ == "__main__": unittest.main()