agent-builder/tests/test_template_init.py

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("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)
                ]
            }
        ]