test: add 14 unit tests for utils/helper.py, utils/telemetry.py (#33)

* Add unit tests for utils/helper.py,utils/telemetry.py

* Update test_helpers.py

* Update test_helpers.py

---------

Co-authored-by: kthota-g <[email protected]>

Author: mindpower

tests/utils/test_helpers.py

@@ -0,0 +1,154 @@
+from typing import Any
+import pytest
+from unittest.mock import MagicMock
+from uuid import uuid4
+from a2a.utils.helpers import (
+    create_task_obj,
+    append_artifact_to_task,
+    build_text_artifact,
+    validate,
+)
+from a2a.types import (
+    Artifact,
+    MessageSendParams,
+    Message,
+    Task,
+    TaskArtifactUpdateEvent,
+    TaskState,
+    TaskStatus,
+    TextPart,
+    Part,
+)
+from a2a.utils.errors import ServerError, UnsupportedOperationError
+
+# --- Helper Data ---
+TEXT_PART_DATA: dict[str, Any] = {'type': 'text', 'text': 'Hello'}
+
+MINIMAL_MESSAGE_USER: dict[str, Any] = {
+    'role': 'user',
+    'parts': [TEXT_PART_DATA],
+    'messageId': 'msg-123',
+    'type': 'message',
+}
+
+MINIMAL_TASK_STATUS: dict[str, Any] = {'state': 'submitted'}
+
+MINIMAL_TASK: dict[str, Any] = {
+    'id': 'task-abc',
+    'contextId': 'session-xyz',
+    'status': MINIMAL_TASK_STATUS,
+    'type': 'task',
+}
+
+# Test create_task_obj
+def test_create_task_obj():
+    message = Message(**MINIMAL_MESSAGE_USER)
+    send_params = MessageSendParams(message=message)
+
+    task = create_task_obj(send_params)
+    assert task.id is not None
+    assert task.contextId == message.contextId
+    assert task.status.state == TaskState.submitted
+    assert len(task.history) == 1
+    assert task.history[0] == message
+
+
+# Test append_artifact_to_task
+def test_append_artifact_to_task():
+     # Prepare base task   
+    task = Task(**MINIMAL_TASK)
+    assert task.id == 'task-abc'
+    assert task.contextId == 'session-xyz'
+    assert task.status.state == TaskState.submitted
+    assert task.history is None
+    assert task.artifacts is None
+    assert task.metadata is None
+
+    # Prepare appending artifact and event
+    artifact_1 = Artifact(
+        artifactId="artifact-123", parts=[Part(root=TextPart(text="Hello"))]
+    )
+    append_event_1 = TaskArtifactUpdateEvent(artifact=artifact_1, append=False, taskId="123", contextId="123")
+
+    # Test adding a new artifact (not appending)
+    append_artifact_to_task(task, append_event_1)
+    assert len(task.artifacts) == 1
+    assert task.artifacts[0].artifactId == "artifact-123"
+    assert task.artifacts[0].name == None
+    assert len(task.artifacts[0].parts) == 1
+    assert task.artifacts[0].parts[0].root.text == "Hello"
+
+    # Test replacing the artifact
+    artifact_2 = Artifact(
+        artifactId="artifact-123", name="updated name", parts=[Part(root=TextPart(text="Updated"))]
+    )
+    append_event_2 = TaskArtifactUpdateEvent(artifact=artifact_2, append=False, taskId="123", contextId="123")
+    append_artifact_to_task(task, append_event_2)
+    assert len(task.artifacts) == 1  # Should still have one artifact
+    assert task.artifacts[0].artifactId == "artifact-123"
+    assert task.artifacts[0].name == "updated name"
+    assert len(task.artifacts[0].parts) == 1
+    assert task.artifacts[0].parts[0].root.text == "Updated"
+
+    # Test appending parts to an existing artifact
+    artifact_with_parts = Artifact(
+        artifactId="artifact-123", parts=[Part(root=TextPart(text="Part 2"))]
+    )
+    append_event_3 = TaskArtifactUpdateEvent(artifact=artifact_with_parts, append=True, taskId="123", contextId="123")
+    append_artifact_to_task(task, append_event_3)
+    assert len(task.artifacts[0].parts) == 2
+    assert task.artifacts[0].parts[0].root.text == "Updated"
+    assert task.artifacts[0].parts[1].root.text == "Part 2"
+
+    # Test adding another new artifact
+    another_artifact_with_parts = Artifact(
+        artifactId="new_artifact", parts=[Part(root=TextPart(text="new artifact Part 1"))]
+    )
+    append_event_4 = TaskArtifactUpdateEvent(artifact=another_artifact_with_parts, append=False, taskId="123", contextId="123")
+    append_artifact_to_task(task, append_event_4)
+    assert len(task.artifacts) == 2
+    assert task.artifacts[0].artifactId == "artifact-123"
+    assert task.artifacts[1].artifactId == "new_artifact"
+    assert len(task.artifacts[0].parts) == 2
+    assert len(task.artifacts[1].parts) == 1
+
+    # Test appending part to a task that does not have a matching artifact
+    non_existing_artifact_with_parts = Artifact(
+        artifactId="artifact-456", parts=[Part(root=TextPart(text="Part 1"))]
+    )
+    append_event_5 = TaskArtifactUpdateEvent(artifact=non_existing_artifact_with_parts, append=True, taskId="123", contextId="123")
+    append_artifact_to_task(task, append_event_5)
+    assert len(task.artifacts) == 2
+    assert len(task.artifacts[0].parts) == 2
+    assert len(task.artifacts[1].parts) == 1
+
+# Test build_text_artifact
+def test_build_text_artifact():
+    artifact_id = "text_artifact"
+    text = "This is a sample text"
+    artifact = build_text_artifact(text, artifact_id)
+
+    assert artifact.artifactId == artifact_id
+    assert len(artifact.parts) == 1
+    assert artifact.parts[0].root.text == text
+
+
+# Test validate decorator
+def test_validate_decorator():
+    class TestClass:
+        condition = True
+
+        @validate(lambda self: self.condition, "Condition not met")
+        def test_method(self):
+            return "Success"
+
+    obj = TestClass()
+
+    # Test passing condition
+    assert obj.test_method() == "Success"
+
+    # Test failing condition
+    obj.condition = False
+    with pytest.raises(ServerError) as exc_info:
+        obj.test_method()
+    assert "Condition not met" in str(exc_info.value)

tests/utils/test_telemetry.py

@@ -0,0 +1,154 @@
+import pytest
+import types
+import asyncio
+from unittest import mock
+from a2a.utils.telemetry import trace_function, trace_class
+
+@pytest.fixture
+def mock_span():
+    span = mock.MagicMock()
+    return span
+
+@pytest.fixture
+def mock_tracer(mock_span):
+    tracer = mock.MagicMock()
+    tracer.start_as_current_span.return_value.__enter__.return_value = mock_span
+    tracer.start_as_current_span.return_value.__exit__.return_value = False
+    return tracer
+
+@pytest.fixture(autouse=True)
+def patch_trace_get_tracer(mock_tracer):
+    with mock.patch("opentelemetry.trace.get_tracer", return_value=mock_tracer):
+        yield
+
+def test_trace_function_sync_success(mock_span):
+    @trace_function
+    def foo(x, y):
+        return x + y
+
+    result = foo(2, 3)
+    assert result == 5
+    mock_span.set_status.assert_called()
+    mock_span.set_status.assert_any_call(mock.ANY)
+    mock_span.record_exception.assert_not_called()
+
+def test_trace_function_sync_exception(mock_span):
+    @trace_function
+    def bar():
+        raise ValueError("fail")
+
+    with pytest.raises(ValueError):
+        bar()
+    mock_span.record_exception.assert_called()
+    mock_span.set_status.assert_any_call(mock.ANY, description="fail")
+
+def test_trace_function_sync_attribute_extractor_called(mock_span):
+    called = {}
+    def attr_extractor(span, args, kwargs, result, exception):
+        called['called'] = True
+        assert span is mock_span
+        assert exception is None
+        assert result == 42
+
+    @trace_function(attribute_extractor=attr_extractor)
+    def foo():
+        return 42
+
+    foo()
+    assert called['called']
+
+def test_trace_function_sync_attribute_extractor_error_logged(mock_span):
+    with mock.patch("a2a.utils.telemetry.logger") as logger:
+        def attr_extractor(span, args, kwargs, result, exception):
+            raise RuntimeError("attr fail")
+
+        @trace_function(attribute_extractor=attr_extractor)
+        def foo():
+            return 1
+
+        foo()
+        logger.error.assert_any_call(mock.ANY)
+
+@pytest.mark.asyncio
+async def test_trace_function_async_success(mock_span):
+    @trace_function
+    async def foo(x):
+        await asyncio.sleep(0)
+        return x * 2
+
+    result = await foo(4)
+    assert result == 8
+    mock_span.set_status.assert_called()
+    mock_span.record_exception.assert_not_called()
+
+@pytest.mark.asyncio
+async def test_trace_function_async_exception(mock_span):
+    @trace_function
+    async def bar():
+        await asyncio.sleep(0)
+        raise RuntimeError("async fail")
+
+    with pytest.raises(RuntimeError):
+        await bar()
+    mock_span.record_exception.assert_called()
+    mock_span.set_status.assert_any_call(mock.ANY, description="async fail")
+
+@pytest.mark.asyncio
+async def test_trace_function_async_attribute_extractor_called(mock_span):
+    called = {}
+    def attr_extractor(span, args, kwargs, result, exception):
+        called['called'] = True
+        assert exception is None
+        assert result == 99
+
+    @trace_function(attribute_extractor=attr_extractor)
+    async def foo():
+        return 99
+
+    await foo()
+    assert called['called']
+
+def test_trace_function_with_args_and_attributes(mock_span):
+    @trace_function(span_name="custom.span", attributes={"foo": "bar"})
+    def foo():
+        return 1
+
+    foo()
+    mock_span.set_attribute.assert_any_call("foo", "bar")
+
+def test_trace_class_exclude_list(mock_span):
+    @trace_class(exclude_list=["skip_me"])
+    class MyClass:
+        def a(self): return "a"
+        def skip_me(self): return "skip"
+        def __str__(self): return "str"
+
+    obj = MyClass()
+    assert obj.a() == "a"
+    assert obj.skip_me() == "skip"
+    # Only 'a' is traced, not 'skip_me' or dunder
+    assert hasattr(obj.a, "__wrapped__")
+    assert not hasattr(obj.skip_me, "__wrapped__")
+
+def test_trace_class_include_list(mock_span):
+    @trace_class(include_list=["only_this"])
+    class MyClass:
+        def only_this(self): return "yes"
+        def not_this(self): return "no"
+
+    obj = MyClass()
+    assert obj.only_this() == "yes"
+    assert obj.not_this() == "no"
+    assert hasattr(obj.only_this, "__wrapped__")
+    assert not hasattr(obj.not_this, "__wrapped__")
+
+def test_trace_class_dunder_not_traced(mock_span):
+    @trace_class()
+    class MyClass:
+        def __init__(self): self.x = 1
+        def foo(self): return "foo"
+
+    obj = MyClass()
+    assert obj.foo() == "foo"
+    assert hasattr(obj.foo, "__wrapped__")
+    assert hasattr(obj, "x")