Restore tests to original

Author: gaborfeher

tests/contrib/tasks/test_vertex_task_converter.py

@@ -1,5 +1,3 @@
-"""Tests for vertex_task_converter mappings."""
-
 import base64
 
 import pytest
@@ -8,17 +6,14 @@
 pytest.importorskip(
     'vertexai', reason='Vertex Task Converter tests require vertexai'
 )
-from google.genai import types as genai_types
 from vertexai import types as vertexai_types
-
+from google.genai import types as genai_types
 from a2a.contrib.tasks.vertex_task_converter import (
     to_sdk_artifact,
-    to_sdk_message,
     to_sdk_part,
     to_sdk_task,
     to_sdk_task_state,
     to_stored_artifact,
-    to_stored_message,
     to_stored_part,
     to_stored_task,
     to_stored_task_state,
@@ -29,114 +24,146 @@
     FilePart,
     FileWithBytes,
     FileWithUri,
-    Message,
     Part,
-    Role,
     Task,
     TaskState,
     TaskStatus,
     TextPart,
 )
 
 
-def test_artifact_conversion_symmetry() -> None:
-    """Test converting an Artifact to TaskArtifact and back restores everything."""
-    original_artifact = Artifact(
-        artifact_id='art123',
-        name='My cool artifact',
-        description='A very interesting description',
-        extensions=['ext1', 'ext2'],
-        metadata={'custom': 'value'},
-        parts=[
-            Part(
-                root=TextPart(
-                    text='hello', metadata={'part_meta': 'hello_meta'}
-                )
-            ),
-            Part(root=DataPart(data={'foo': 'bar'})),  # no metadata
-        ],
+def test_to_sdk_task_state() -> None:
+    assert (
+        to_sdk_task_state(vertexai_types.A2aTaskState.STATE_UNSPECIFIED)
+        == TaskState.unknown
+    )
+    assert (
+        to_sdk_task_state(vertexai_types.A2aTaskState.SUBMITTED)
+        == TaskState.submitted
+    )
+    assert (
+        to_sdk_task_state(vertexai_types.A2aTaskState.WORKING)
+        == TaskState.working
+    )
+    assert (
+        to_sdk_task_state(vertexai_types.A2aTaskState.COMPLETED)
+        == TaskState.completed
+    )
+    assert (
+        to_sdk_task_state(vertexai_types.A2aTaskState.CANCELLED)
+        == TaskState.canceled
+    )
+    assert (
+        to_sdk_task_state(vertexai_types.A2aTaskState.FAILED)
+        == TaskState.failed
+    )
+    assert (
+        to_sdk_task_state(vertexai_types.A2aTaskState.REJECTED)
+        == TaskState.rejected
+    )
+    assert (
+        to_sdk_task_state(vertexai_types.A2aTaskState.INPUT_REQUIRED)
+        == TaskState.input_required
+    )
+    assert (
+        to_sdk_task_state(vertexai_types.A2aTaskState.AUTH_REQUIRED)
+        == TaskState.auth_required
     )
+    assert to_sdk_task_state(999) == TaskState.unknown  # type: ignore
 
-    stored = to_stored_artifact(original_artifact)
-    assert isinstance(stored, vertexai_types.TaskArtifact)
-
-    # ensure it was populated correctly
-    assert stored.display_name == 'My cool artifact'
-    assert stored.description == 'A very interesting description'
-    assert stored.metadata['__vertex_compat_v'] == 1.0
-
-    restored_artifact = to_sdk_artifact(stored)
-
-    assert restored_artifact.artifact_id == original_artifact.artifact_id
-    assert restored_artifact.name == original_artifact.name
-    assert restored_artifact.description == original_artifact.description
-    assert restored_artifact.extensions == original_artifact.extensions
-    assert restored_artifact.metadata == original_artifact.metadata
-
-    assert len(restored_artifact.parts) == 2
-    assert isinstance(restored_artifact.parts[0].root, TextPart)
-    assert restored_artifact.parts[0].root.text == 'hello'
-    assert restored_artifact.parts[0].root.metadata == {
-        'part_meta': 'hello_meta'
-    }
-
-    assert isinstance(restored_artifact.parts[1].root, DataPart)
-    assert restored_artifact.parts[1].root.data == {'foo': 'bar'}
-    assert restored_artifact.parts[1].root.metadata is None
-
-
-def test_message_conversion_symmetry() -> None:
-    """Test converting a Message to TaskMessage and back restores everything."""
-    original_message = Message(
-        message_id='msg456',
-        role=Role.agent,
-        context_id='ctx1',
-        task_id='tsk1',
-        reference_task_ids=['tsk2', 'tsk3'],
-        extensions=['ext_msg'],
-        metadata={'msg_meta': 42},
-        parts=[
-            Part(root=TextPart(text='message text')),
-        ],
+
+def test_to_stored_task_state() -> None:
+    assert (
+        to_stored_task_state(TaskState.unknown)
+        == vertexai_types.A2aTaskState.STATE_UNSPECIFIED
+    )
+    assert (
+        to_stored_task_state(TaskState.submitted)
+        == vertexai_types.A2aTaskState.SUBMITTED
+    )
+    assert (
+        to_stored_task_state(TaskState.working)
+        == vertexai_types.A2aTaskState.WORKING
+    )
+    assert (
+        to_stored_task_state(TaskState.completed)
+        == vertexai_types.A2aTaskState.COMPLETED
+    )
+    assert (
+        to_stored_task_state(TaskState.canceled)
+        == vertexai_types.A2aTaskState.CANCELLED
+    )
+    assert (
+        to_stored_task_state(TaskState.failed)
+        == vertexai_types.A2aTaskState.FAILED
+    )
+    assert (
+        to_stored_task_state(TaskState.rejected)
+        == vertexai_types.A2aTaskState.REJECTED
+    )
+    assert (
+        to_stored_task_state(TaskState.input_required)
+        == vertexai_types.A2aTaskState.INPUT_REQUIRED
+    )
+    assert (
+        to_stored_task_state(TaskState.auth_required)
+        == vertexai_types.A2aTaskState.AUTH_REQUIRED
     )
 
-    stored = to_stored_message(original_message)
-    assert stored is not None
-    assert isinstance(stored, vertexai_types.TaskMessage)
 
-    assert stored.message_id == 'msg456'
-    assert stored.role == 'agent'
-    assert stored.metadata['__vertex_compat_v'] == 1.0
+def test_to_stored_part_text() -> None:
+    sdk_part = Part(root=TextPart(text='hello world'))
+    stored_part = to_stored_part(sdk_part)
+    assert stored_part.text == 'hello world'
+    assert not stored_part.inline_data
+    assert not stored_part.file_data
 
-    restored_message = to_sdk_message(stored)
-    assert restored_message is not None
 
-    assert restored_message.message_id == original_message.message_id
-    assert restored_message.role == original_message.role
-    # context_id and task_id are not serialized via Message metadata in Go implementation but via Task,
-    # but reference_task_ids and extensions ARE part of Message metadata.
-    assert (
-        restored_message.reference_task_ids
-        == original_message.reference_task_ids
-    )
-    assert restored_message.extensions == original_message.extensions
-    assert restored_message.metadata == original_message.metadata
+def test_to_stored_part_data() -> None:
+    sdk_part = Part(root=DataPart(data={'key': 'value'}))
+    stored_part = to_stored_part(sdk_part)
+    assert stored_part.inline_data is not None
+    assert stored_part.inline_data.mime_type == 'application/json'
+    assert stored_part.inline_data.data == b'{"key": "value"}'
 
-    assert len(restored_message.parts) == 1
-    assert isinstance(restored_message.parts[0].root, TextPart)
-    assert restored_message.parts[0].root.text == 'message text'
-    assert restored_message.parts[0].root.metadata is None
 
+def test_to_stored_part_file_bytes() -> None:
+    encoded_b64 = base64.b64encode(b'test data').decode('utf-8')
+    sdk_part = Part(
+        root=FilePart(
+            file=FileWithBytes(
+                bytes=encoded_b64,
+                mime_type='text/plain',
+            )
+        )
+    )
+    stored_part = to_stored_part(sdk_part)
+    assert stored_part.inline_data is not None
+    assert stored_part.inline_data.mime_type == 'text/plain'
+    assert stored_part.inline_data.data == b'test data'
 
-def test_to_stored_part_unsupported() -> None:
-    part = Part.model_construct(
-        root=Task(  # type: ignore[arg-type]
-            id='invalid-part',
-            context_id='ctx',
-            status=TaskStatus(state=TaskState.submitted),
-            history=[],
+
+def test_to_stored_part_file_uri() -> None:
+    sdk_part = Part(
+        root=FilePart(
+            file=FileWithUri(
+                uri='gs://test-bucket/file.txt',
+                mime_type='text/plain',
+            )
         )
     )
+    stored_part = to_stored_part(sdk_part)
+    assert stored_part.file_data is not None
+    assert stored_part.file_data.mime_type == 'text/plain'
+    assert stored_part.file_data.file_uri == 'gs://test-bucket/file.txt'
+
+
+def test_to_stored_part_unsupported() -> None:
+    class BadPart:
+        pass
+
+    part = Part(root=TextPart(text='t'))
+    part.root = BadPart()  # type: ignore
     with pytest.raises(ValueError, match='Unsupported part type'):
         to_stored_part(part)