124 lines
4.2 KiB
Python
124 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import tarfile
|
|
import unittest
|
|
|
|
from agent_builder.tools import (
|
|
_BUILDER_STATE_FILE,
|
|
_deployment_drift_error,
|
|
_render_a2a_init_template,
|
|
_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("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_stale_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_drift")
|
|
self.assertEqual(err["current_head_sha"], "abc1234")
|
|
self.assertIsNone(err["workspace_base_head_sha"])
|
|
|
|
def test_deployment_drift_error_allows_matching_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": "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"}',
|
|
})
|
|
|
|
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 _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])}
|
|
|
|
|
|
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)
|
|
]
|
|
}
|
|
]
|