File size: 2,650 Bytes
54b0fbc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
"""Tests for AgentCore runtime URL parsing and response mapping."""

from __future__ import annotations

import json

from pi_test_support import ensure_agent_redact_paths

ensure_agent_redact_paths()

from agentcore_runtime import (  # noqa: E402
    AgentCoreAgentRuntime,
    agentcore_runtime_url,
    parse_agentcore_runtime_url,
)


def test_parse_agentcore_runtime_url_from_base():
    url = (
        "https://bedrock-agentcore.eu-west-2.amazonaws.com/runtimes/"
        "arn%3Aaws%3Abedrock-agentcore%3Aeu-west-2%3A404053085091%3Aruntime%2FRedactionAgent"
    )
    region, arn = parse_agentcore_runtime_url(url)
    assert region == "eu-west-2"
    assert (
        arn == "arn:aws:bedrock-agentcore:eu-west-2:404053085091:runtime/RedactionAgent"
    )


def test_parse_agentcore_runtime_url_strips_invocations_suffix():
    url = (
        "https://bedrock-agentcore.eu-west-2.amazonaws.com/runtimes/"
        "arn%3Aaws%3Abedrock-agentcore%3Aeu-west-2%3A404053085091%3Aruntime%2FRedactionAgent"
        "/invocations"
    )
    region, arn = parse_agentcore_runtime_url(url)
    assert region == "eu-west-2"
    assert arn.endswith("runtime/RedactionAgent")


def test_agentcore_runtime_url_strips_invocations(monkeypatch):
    monkeypatch.setenv(
        "AGENTCORE_RUNTIME_URL",
        "https://bedrock-agentcore.eu-west-2.amazonaws.com/runtimes/arn%3Ax/invocations",
    )
    assert agentcore_runtime_url().endswith("arn%3Ax")
    assert not agentcore_runtime_url().endswith("/invocations")


def test_iter_json_response_result_field():
    runtime = AgentCoreAgentRuntime()
    events = list(runtime._iter_json_response(json.dumps({"result": "hello"}).encode()))
    assert len(events) == 1
    assert events[0].kind == "text_snapshot"
    assert events[0].text == "hello"


def test_map_message_update_tool_calls():
    runtime = AgentCoreAgentRuntime(session_hash="sess")
    event = {
        "type": "message_update",
        "role": "assistant",
        "content": "Running doc_redact.",
        "tool_calls": [{"name": "doc_redact", "args": {"pdf_relative_path": "a.pdf"}}],
    }
    kinds = [e.kind for e in runtime._map_agentcore_event(event)]
    assert kinds == ["tool_start", "text_snapshot"]


def test_map_message_update_tool_result():
    runtime = AgentCoreAgentRuntime(session_hash="sess")
    event = {
        "type": "message_update",
        "role": "tool",
        "tool_name": "doc_redact",
        "content": '{"message": "done"}',
    }
    events = list(runtime._map_agentcore_event(event))
    assert len(events) == 1
    assert events[0].kind == "tool_end"
    assert events[0].tool_name == "doc_redact"