from __future__ import annotations import io import tarfile import unittest from agent_builder.tools import ( _BUILDER_STATE_FILE, _BUILDER_INTERNAL_PREFIX, _deployment_drift_error, _files_from_tarball, _parse_supporting_skill_files, _replace_workspace_files, _render_a2a_init_template, _render_skill_md, _tarball_workspace_dir, ) class TemplateInitTests(unittest.TestCase): def test_render_a2a_init_template_matches_sdk_starter(self) -> None: files = _render_a2a_init_template( "research-agent", description="Research helper", ) self.assertEqual(set(files), {"agent.py", "a2a.yaml", "requirements.txt"}) self.assertIn('name = "research-agent"', files["agent.py"]) self.assertIn('description = "Research helper"', files["agent.py"]) self.assertIn("LLMProvisioning.CALLER_PROVIDED", files["agent.py"]) self.assertIn("WorkspaceAccess.dynamic", files["agent.py"]) self.assertIn("RUNTIME_SKILLS_ROOT", files["agent.py"]) self.assertIn("DEEPAGENTS_RECURSION_LIMIT = 500", files["agent.py"]) self.assertIn( 'config={"recursion_limit": DEEPAGENTS_RECURSION_LIMIT}', files["agent.py"], ) self.assertIn("skills=skill_sources or None", files["agent.py"]) self.assertIn("ctx.llm", files["agent.py"]) self.assertIn("ctx.workspace_backend()", files["agent.py"]) self.assertIn("create_deep_agent", files["agent.py"]) self.assertIn("name: research-agent", files["a2a.yaml"]) self.assertIn("entrypoint: agent:ResearchAgent", files["a2a.yaml"]) self.assertIn("deepagents>=0.5.0", files["requirements.txt"]) def test_render_a2a_init_template_rejects_invalid_names(self) -> None: with self.assertRaises(ValueError): _render_a2a_init_template("Bad Name") def test_agent_projects_can_diverge_from_template(self) -> None: files = _render_a2a_init_template("regex-explainer") files["agent.py"] = "from a2a_pack import A2AAgent\n" self.assertIn("A2AAgent", files["agent.py"]) def test_deployment_drift_error_blocks_untracked_existing_workspace(self) -> None: latest = {"deploy_id": "dpl_1", "head_sha": "abc1234"} err = _deployment_drift_error("demo-agent", latest, {}, force=False) self.assertIsNotNone(err) assert err is not None self.assertEqual(err["error"], "workspace_untracked") self.assertEqual(err["current_head_sha"], "abc1234") def test_deployment_drift_error_allows_tracked_or_forced_workspace(self) -> None: latest = {"deploy_id": "dpl_1", "head_sha": "abc1234"} self.assertIsNone( _deployment_drift_error( "demo-agent", latest, {"repo_head_sha": "abc1234"}, force=False, ) ) self.assertIsNone( _deployment_drift_error( "demo-agent", latest, {"repo_head_sha": "repohead999"}, force=False, ) ) self.assertIsNone( _deployment_drift_error( "demo-agent", latest, {"repo_head_sha": "old9999"}, force=True, ) ) def test_tarball_workspace_excludes_builder_state(self) -> None: prefix = "agents/demo-agent/" s3 = _FakeS3({ prefix + "agent.py": b"print('ok')\n", prefix + _BUILDER_STATE_FILE: b'{"repo_head_sha":"abc1234"}', prefix + _BUILDER_INTERNAL_PREFIX + "skills/x/SKILL.md": b"hidden", }) bundle = _tarball_workspace_dir(s3, "bucket", prefix) with tarfile.open(fileobj=io.BytesIO(bundle), mode="r:gz") as tf: self.assertEqual(tf.getnames(), ["agent.py"]) def test_render_skill_md_creates_valid_frontmatter(self) -> None: text = _render_skill_md( "market-research", "Plan market research and synthesize findings.", "Use subagents for separate market segments.", ) self.assertIn("name: market-research", text) self.assertIn('description: "Plan market research', text) self.assertIn("# market-research", text) def test_parse_supporting_skill_files_rejects_unsafe_paths(self) -> None: with self.assertRaises(ValueError): _parse_supporting_skill_files('{"../secrets.txt": "bad"}') files = _parse_supporting_skill_files( '{"references/schema.md": "schema", "scripts/check.py": "print(1)"}' ) self.assertEqual( sorted(files), ["references/schema.md", "scripts/check.py"], ) def test_files_from_tarball_rejects_unsafe_paths(self) -> None: bundle = _tarball({"../agent.py": "bad"}) with self.assertRaises(ValueError): _files_from_tarball(bundle) def test_replace_workspace_files_deletes_stale_objects(self) -> None: prefix = "agents/demo-agent/" s3 = _FakeS3({ prefix + "old.py": b"old", prefix + _BUILDER_STATE_FILE: b"{}", }) written = _replace_workspace_files( s3, "bucket", prefix, {"agent.py": b"print('ok')\n", "a2a.yaml": b"name: demo-agent\n"}, ) self.assertEqual( sorted(s3.objects), [prefix + "a2a.yaml", prefix + "agent.py"], ) self.assertEqual( written, [ {"path": "a2a.yaml", "size": len(b"name: demo-agent\n")}, {"path": "agent.py", "size": len(b"print('ok')\n")}, ], ) def _tarball(files: dict[str, str]) -> bytes: buf = io.BytesIO() with tarfile.open(fileobj=buf, mode="w:gz") as tf: for path, text in files.items(): body = text.encode("utf-8") info = tarfile.TarInfo(path) info.size = len(body) tf.addfile(info, io.BytesIO(body)) return buf.getvalue() class _FakeS3: def __init__(self, objects: dict[str, bytes]) -> None: self.objects = objects def get_paginator(self, name: str) -> "_FakePaginator": assert name == "list_objects_v2" return _FakePaginator(self.objects) def get_object(self, *, Bucket: str, Key: str) -> dict[str, io.BytesIO]: return {"Body": io.BytesIO(self.objects[Key])} def put_object( self, *, Bucket: str, Key: str, Body: bytes, ContentType: str, ) -> None: self.objects[Key] = bytes(Body) def delete_objects(self, *, Bucket: str, Delete: dict[str, object]) -> None: for obj in Delete["Objects"]: # type: ignore[index] self.objects.pop(obj["Key"], None) class _FakePaginator: def __init__(self, objects: dict[str, bytes]) -> None: self.objects = objects def paginate(self, *, Bucket: str, Prefix: str) -> list[dict[str, object]]: return [ { "Contents": [ {"Key": key, "Size": len(value)} for key, value in sorted(self.objects.items()) if key.startswith(Prefix) ] } ]