fix: consistent params order & tests fix

Author: martimfasantos

src/a2a/helpers/proto_helpers.py

@@ -27,9 +27,9 @@
 
 def new_message(
     parts: list[Part],
-    role: Role = Role.ROLE_AGENT,
     context_id: str | None = None,
     task_id: str | None = None,
+    role: Role = Role.ROLE_AGENT,
 ) -> Message:
     """Creates a new message containing a list of Parts."""
     return Message(
@@ -44,16 +44,16 @@ def new_message(
 def new_text_message(
     text: str,
     media_type: str | None = None,
-    role: Role = Role.ROLE_AGENT,
     context_id: str | None = None,
     task_id: str | None = None,
+    role: Role = Role.ROLE_AGENT,
 ) -> Message:
     """Creates a new message containing a single text Part."""
     return new_message(
         parts=[new_text_part(text, media_type=media_type)],
-        role=role,
-        task_id=task_id,
         context_id=context_id,
+        task_id=task_id,
+        role=role,
     )
 
 
@@ -65,85 +65,85 @@ def get_message_text(message: Message, delimiter: str = '\n') -> str:
 def new_data_message(
     data: Any,
     media_type: str | None = None,
-    role: Role = Role.ROLE_AGENT,
     context_id: str | None = None,
     task_id: str | None = None,
+    role: Role = Role.ROLE_AGENT,
 ) -> Message:
     """Creates a new message containing a single data Part.
 
     Args:
         data: JSON-serializable data to embed (dict, list, str, etc.).
         media_type: Optional MIME type of the part content (e.g., "text/plain", "application/json", "image/png").
-        role: The role of the message sender (default: ROLE_AGENT).
         context_id: Optional context ID.
         task_id: Optional task ID.
+        role: The role of the message sender (default: ROLE_AGENT).
 
     Returns:
         A Message with a single data Part.
     """
     return new_message(
         parts=[new_data_part(data, media_type=media_type)],
-        role=role,
         context_id=context_id,
         task_id=task_id,
+        role=role,
     )
 
 
 def new_raw_message(  # noqa: PLR0913
     raw: bytes,
     media_type: str | None = None,
     filename: str | None = None,
-    role: Role = Role.ROLE_AGENT,
     context_id: str | None = None,
     task_id: str | None = None,
+    role: Role = Role.ROLE_AGENT,
 ) -> Message:
     """Creates a new message containing a single raw bytes Part.
 
     Args:
         raw: The raw bytes content.
         media_type: Optional MIME type (e.g. 'image/png').
         filename: Optional filename.
-        role: The role of the message sender (default: ROLE_AGENT).
         context_id: Optional context ID.
         task_id: Optional task ID.
+        role: The role of the message sender (default: ROLE_AGENT).
 
     Returns:
         A Message with a single raw Part.
     """
     return new_message(
         parts=[new_raw_part(raw, media_type=media_type, filename=filename)],
-        role=role,
         context_id=context_id,
         task_id=task_id,
+        role=role,
     )
 
 
 def new_url_message(  # noqa: PLR0913
     url: str,
     media_type: str | None = None,
     filename: str | None = None,
-    role: Role = Role.ROLE_AGENT,
     context_id: str | None = None,
     task_id: str | None = None,
+    role: Role = Role.ROLE_AGENT,
 ) -> Message:
     """Creates a new message containing a single URL Part.
 
     Args:
         url: The URL pointing to the file content.
         media_type: Optional MIME type (e.g. 'image/png').
         filename: Optional filename.
-        role: The role of the message sender (default: ROLE_AGENT).
         context_id: Optional context ID.
         task_id: Optional task ID.
+        role: The role of the message sender (default: ROLE_AGENT).
 
     Returns:
         A Message with a single URL Part.
     """
     return new_message(
         parts=[new_url_part(url, media_type=media_type, filename=filename)],
-        role=role,
         context_id=context_id,
         task_id=task_id,
+        role=role,
     )
 
 

tests/helpers/test_proto_helpers.py

@@ -1,47 +1,49 @@
 """Tests for proto helpers."""
 
 import pytest
+
 from a2a.helpers.proto_helpers import (
-    new_message,
-    new_text_message,
-    new_data_message,
-    new_raw_message,
-    new_url_message,
+    get_artifact_text,
     get_message_text,
+    get_stream_response_text,
+    get_text_parts,
     new_artifact,
-    new_text_artifact,
     new_data_artifact,
+    new_data_message,
+    new_data_part,
+    new_message,
     new_raw_artifact,
-    new_url_artifact,
-    get_artifact_text,
-    new_task_from_user_message,
+    new_raw_message,
+    new_raw_part,
     new_task,
-    get_text_parts,
+    new_task_from_user_message,
+    new_text_artifact,
+    new_text_artifact_update_event,
+    new_text_message,
     new_text_part,
-    new_data_part,
-    new_raw_part,
-    new_url_part,
     new_text_status_update_event,
-    new_text_artifact_update_event,
-    get_stream_response_text,
+    new_url_artifact,
+    new_url_message,
+    new_url_part,
 )
 from a2a.types.a2a_pb2 import (
+    Artifact,
+    Message,
     Part,
     Role,
-    Message,
-    Artifact,
+    StreamResponse,
     Task,
     TaskState,
-    StreamResponse,
 )
 
+
 # --- Message Helpers Tests ---
 
 
 def test_new_message() -> None:
     parts = [Part(text='hello')]
     msg = new_message(
-        parts=parts, role=Role.ROLE_USER, context_id='ctx1', task_id='task1'
+        parts, context_id='ctx1', task_id='task1', role=Role.ROLE_USER
     )
     assert msg.role == Role.ROLE_USER
     assert msg.parts == parts
@@ -52,11 +54,16 @@ def test_new_message() -> None:
 
 def test_new_text_message() -> None:
     msg = new_text_message(
-        text='hello', context_id='ctx1', task_id='task1', role=Role.ROLE_USER
+        'hello',
+        media_type='text/plain',
+        context_id='ctx1',
+        task_id='task1',
+        role=Role.ROLE_USER,
     )
     assert msg.role == Role.ROLE_USER
     assert len(msg.parts) == 1
     assert msg.parts[0].text == 'hello'
+    assert msg.parts[0].media_type == 'text/plain'
     assert msg.context_id == 'ctx1'
     assert msg.task_id == 'task1'
     assert msg.message_id != ''
@@ -65,6 +72,7 @@ def test_new_text_message() -> None:
 def test_new_data_message() -> None:
     msg = new_data_message(
         data={'key': 'value'},
+        media_type='application/json',
         context_id='ctx1',
         task_id='task1',
         role=Role.ROLE_USER,
@@ -73,6 +81,7 @@ def test_new_data_message() -> None:
     assert len(msg.parts) == 1
     assert msg.parts[0].HasField('data')
     assert msg.parts[0].data.struct_value.fields['key'].string_value == 'value'
+    assert msg.parts[0].media_type == 'application/json'
     assert msg.context_id == 'ctx1'
     assert msg.task_id == 'task1'
     assert msg.message_id != ''
@@ -85,7 +94,9 @@ def test_new_raw_message() -> None:
         filename='img.png',
         context_id='ctx1',
         task_id='task1',
+        role=Role.ROLE_USER,
     )
+    assert msg.role == Role.ROLE_USER
     assert len(msg.parts) == 1
     assert msg.parts[0].HasField('raw')
     assert msg.parts[0].raw == b'\x89PNG'
@@ -103,7 +114,9 @@ def test_new_url_message() -> None:
         filename='file.pdf',
         context_id='ctx1',
         task_id='task1',
+        role=Role.ROLE_USER,
     )
+    assert msg.role == Role.ROLE_USER
     assert len(msg.parts) == 1
     assert msg.parts[0].HasField('url')
     assert msg.parts[0].url == 'https://example.com/file.pdf'