python(feat): Add FileAttachmentsMixin to TestStep. (#466)

Author: ian-sift

python/lib/sift_client/_tests/sift_types/test_results.py

@@ -2,8 +2,9 @@
 
 from __future__ import annotations
 
+import tempfile
 from datetime import datetime, timedelta, timezone
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, call
 
 import pytest
 
@@ -230,3 +231,41 @@ def test_attachments_property_fetches_files(self, mock_test_report, mock_client)
 
         # Verify result
         assert result == mock_remote_files
+
+    def test_upload_attachment(self, mock_test_report, mock_test_step, mock_client):
+        """Ensure test report and step have FileAttachmentsMixin and it is called correctly."""
+        # Create mock file attachment to be returned
+        mock_file_attachment = MagicMock()
+        mock_file_attachment.description = "Test upload to test report"
+        mock_file_attachment.entity_id = mock_test_report.id_
+        mock_client.file_attachments.upload.return_value = mock_file_attachment
+
+        # Create a temporary test file
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as tmp:
+            tmp.write("Test file content\n")
+            tmp_path = tmp.name
+
+        _ = mock_test_report.upload_attachment(
+            path=tmp_path, description="Test upload to test report"
+        )
+        _ = mock_test_step.upload_attachment(path=tmp_path, description="Test upload to test step")
+
+        # Verify file_attachments.upload was called with correct parameters
+        mock_client.file_attachments.upload.assert_has_calls(
+            [
+                call(
+                    path=tmp_path,
+                    entity=mock_test_report,
+                    metadata=None,
+                    description="Test upload to test report",
+                    organization_id=None,
+                ),
+                call(
+                    path=tmp_path,
+                    entity=mock_test_step,
+                    metadata=None,
+                    description="Test upload to test step",
+                    organization_id=None,
+                ),
+            ]
+        )

python/lib/sift_client/resources/file_attachments.py

@@ -19,7 +19,7 @@
         RemoteFileEntityType,
     )
     from sift_client.sift_types.run import Run
-    from sift_client.sift_types.test_report import TestReport
+    from sift_client.sift_types.test_report import TestReport, TestStep
 
 
 class FileAttachmentsAPIAsync(ResourceBase):
@@ -94,7 +94,7 @@ async def list_(
         # metadata TODO: Add to backend
         # metadata: list[Any] | None = None,
         # file specific
-        entities: list[Run | Asset | TestReport] | None = None,
+        entities: list[Run | Asset | TestReport | TestStep] | None = None,
         entity_type: RemoteFileEntityType | None = None,
         entity_ids: list[str] | None = None,
         # common filters
@@ -255,7 +255,7 @@ async def upload(
         self,
         *,
         path: str | Path,
-        entity: Asset | Run | TestReport,
+        entity: Asset | Run | TestReport | TestStep,
         metadata: dict[str, Any] | None = None,
         description: str | None = None,
         organization_id: str | None = None,

python/lib/sift_client/resources/sync_stubs/__init__.pyi

@@ -582,7 +582,7 @@ class FileAttachmentsAPI:
         name_contains: str | None = None,
         name_regex: str | re.Pattern | None = None,
         remote_file_ids: list[str] | None = None,
-        entities: list[Run | Asset | TestReport] | None = None,
+        entities: list[Run | Asset | TestReport | TestStep] | None = None,
         entity_type: RemoteFileEntityType | None = None,
         entity_ids: list[str] | None = None,
         description_contains: str | None = None,
@@ -626,7 +626,7 @@ class FileAttachmentsAPI:
         self,
         *,
         path: str | Path,
-        entity: Asset | Run | TestReport,
+        entity: Asset | Run | TestReport | TestStep,
         metadata: dict[str, Any] | None = None,
         description: str | None = None,
         organization_id: str | None = None,

python/lib/sift_client/sift_types/_base.py

@@ -9,6 +9,7 @@
     Callable,
     ClassVar,
     Generic,
+    Protocol,
     TypeVar,
 )
 
@@ -23,6 +24,19 @@
 SelfT = TypeVar("SelfT", bound="BaseType")
 
 
+class BaseTypeProtocol(Protocol):
+    """Protocol for defining public properties for types that inherit from BaseType."""
+
+    @property
+    def client(self) -> SiftClient: ...
+
+    @property
+    def id_(self) -> str | None: ...
+
+    @property
+    def _id_or_error(self) -> str: ...
+
+
 class BaseType(BaseModel, Generic[ProtoT, SelfT], ABC):
     model_config = ConfigDict(frozen=True)
 

python/lib/sift_client/sift_types/_mixins/file_attachments.py

@@ -1,44 +1,40 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, ClassVar, Protocol
+from typing import TYPE_CHECKING, Any, ClassVar
 
 if TYPE_CHECKING:
     from pathlib import Path
 
-    from sift_client.client import SiftClient
+    from sift_client.sift_types._base import BaseTypeProtocol
     from sift_client.sift_types.file_attachment import FileAttachment
 
 
-class _SupportsFileAttachments(Protocol):
-    """Protocol for types that support file attachments."""
-
-    @property
-    def client(self) -> SiftClient: ...
-
-    @property
-    def id_(self) -> str | None: ...
-
-
 class FileAttachmentsMixin:
     """Mixin for sift_types that support file attachments (remote files).
 
     This mixin assumes the class also inherits from BaseType, which provides:
     - id_: str | None
     - client: SiftClient property
-
-    The entity type is automatically determined from the class name:
-    - Asset -> assets
-    - Run -> runs
-    - TestReport -> test_reports
     """
 
     # Mapping of class names to entity types (REST API format)
     _ENTITY_TYPE_MAP: ClassVar[dict[str, str]] = {
         "Asset": "assets",
         "Run": "runs",
         "TestReport": "test_reports",
+        "TestStep": "test_steps",
     }
 
+    @staticmethod
+    def check_is_supported_entity_type(cls):
+        """Check if the entity type is supported for file attachments.
+
+        Returns:
+            True if the entity type is supported, False otherwise.
+        """
+        if not cls.__class__.__name__ in FileAttachmentsMixin._ENTITY_TYPE_MAP:
+            raise ValueError(f"{cls.__name__} does not support file attachments")
+
     def _get_entity_type_name(self) -> str:
         """Get the entity type string.
 
@@ -49,7 +45,7 @@ def _get_entity_type_name(self) -> str:
             ValueError: If the class name is not in the entity type mapping.
         """
         class_name = self.__class__.__name__
-        entity_type = self._ENTITY_TYPE_MAP.get(class_name)
+        entity_type = FileAttachmentsMixin._ENTITY_TYPE_MAP.get(self.__class__.__name__)
 
         if not entity_type:
             raise ValueError(
@@ -60,24 +56,19 @@ def _get_entity_type_name(self) -> str:
         return entity_type
 
     @property
-    def attachments(self: _SupportsFileAttachments) -> list[FileAttachment]:
+    def attachments(self: BaseTypeProtocol) -> list[FileAttachment]:
         """Get all file attachments for this entity.
 
         Returns:
             A list of FileAttachments associated with this entity.
         """
-        from sift_client.sift_types.asset import Asset
-        from sift_client.sift_types.run import Run
-        from sift_client.sift_types.test_report import TestReport
-
-        if not isinstance(self, (Asset, Run, TestReport)):
-            raise ValueError("Entity is not a valid entity type")
+        FileAttachmentsMixin.check_is_supported_entity_type(self)
         return self.client.file_attachments.list_(
-            entities=[self],
+            entities=[self],  # type: ignore
         )
 
     def delete_attachment(
-        self: _SupportsFileAttachments,
+        self: BaseTypeProtocol,
         file_attachment: list[FileAttachment | str] | FileAttachment | str,
     ) -> None:
         """Delete one or more file attachments.
@@ -88,7 +79,7 @@ def delete_attachment(
         self.client.file_attachments.delete(file_attachments=file_attachment)
 
     def upload_attachment(
-        self: _SupportsFileAttachments,
+        self: BaseTypeProtocol,
         path: str | Path,
         metadata: dict[str, Any] | None = None,
         description: str | None = None,
@@ -105,15 +96,10 @@ def upload_attachment(
         Returns:
             The uploaded FileAttachment.
         """
-        from sift_client.sift_types.asset import Asset
-        from sift_client.sift_types.run import Run
-        from sift_client.sift_types.test_report import TestReport
-
-        if not isinstance(self, (Asset, Run, TestReport)):
-            raise ValueError("Entity is not a valid entity type")
+        FileAttachmentsMixin.check_is_supported_entity_type(self)
         return self.client.file_attachments.upload(
             path=path,
-            entity=self,
+            entity=self,  # type: ignore
             metadata=metadata,
             description=description,
             organization_id=organization_id,

python/lib/sift_client/sift_types/test_report.py

@@ -136,7 +136,7 @@ def to_proto(self) -> TestStepProto:
         return proto
 
 
-class TestStep(BaseType[TestStepProto, "TestStep"]):
+class TestStep(BaseType[TestStepProto, "TestStep"], FileAttachmentsMixin):
     """TestStep model representing a step in a test."""
 
     test_report_id: str