a2a-python/tests/extensions/test_common.py at 3a4f9b958c1a59aaa49d396c7f1db3cf4f3e49c0 · mikeas1/a2a-python · GitHub

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
from a2a.extensions.common import (
    find_extension_by_uri,
    get_requested_extensions,
)
from a2a.types import AgentCapabilities, AgentCard, AgentExtension


def test_get_requested_extensions():
    assert get_requested_extensions([]) == set()
    assert get_requested_extensions(['foo']) == {'foo'}
    assert get_requested_extensions(['foo', 'bar']) == {'foo', 'bar'}
    assert get_requested_extensions(['foo, bar']) == {'foo', 'bar'}
    assert get_requested_extensions(['foo,bar']) == {'foo', 'bar'}
    assert get_requested_extensions(['foo', 'bar,baz']) == {'foo', 'bar', 'baz'}
    assert get_requested_extensions(['foo,, bar', 'baz']) == {
        'foo',
        'bar',
        'baz',
    }
    assert get_requested_extensions([' foo , bar ', 'baz']) == {
        'foo',
        'bar',
        'baz',
    }


def test_find_extension_by_uri():
    ext1 = AgentExtension(uri='foo', description='The Foo extension')
    ext2 = AgentExtension(uri='bar', description='The Bar extension')
    card = AgentCard(
        name='Test Agent',
        description='Test Agent Description',
        version='1.0',
        url='http://test.com',
        skills=[],
        default_input_modes=['text/plain'],
        default_output_modes=['text/plain'],
        capabilities=AgentCapabilities(extensions=[ext1, ext2]),
    )

    assert find_extension_by_uri(card, 'foo') == ext1
    assert find_extension_by_uri(card, 'bar') == ext2
    assert find_extension_by_uri(card, 'baz') is None


def test_find_extension_by_uri_no_extensions():
    card = AgentCard(
        name='Test Agent',
        description='Test Agent Description',
        version='1.0',
        url='http://test.com',
        skills=[],
        default_input_modes=['text/plain'],
        default_output_modes=['text/plain'],
        capabilities=AgentCapabilities(extensions=None),
    )

    assert find_extension_by_uri(card, 'foo') is None