32211 Add state parent class to subscription patch request

src/omnia_timeseries/models.py

@@ -110,14 +110,17 @@ class IMSMetadataModel(TypedDict):
     maxTimeInterval: str
     fieldId: str
 
+
 class IMSMetadataItemsModel(TypedDict):
     items: List[IMSMetadataModel]
 
+
 class GetIMSMetadataResponseModel(TypedDict):
     data: IMSMetadataItemsModel
     count: Optional[int]
     continuationToken: Optional[str]
 
+
 class TimeseriesModel(TypedDict):
     id: str
     name: str
@@ -142,14 +145,31 @@ class GetTimeseriesResponseModel(TypedDict):
     count: Optional[int]
     continuationToken: Optional[str]
 
-class SubscriptionPatchRequestItem(TypedDict, total=False):
+
+class SubscriptionStatePatchModel(TypedDict, total=False):
+    tsdbBackfilled: Optional[bool]
+    tsdbEventsExported: Optional[int]
+    tsdbFirstExportedTime: Optional[str]
+    tsdbLastExportedTime: Optional[str]
+    dlBackfilled: Optional[bool]
+    dlEventsExported: Optional[int]
+    dlFirstExportedTime: Optional[str]
+    dlLastExportedTime: Optional[str]
+    dlParquetFirstExportedTime: Optional[str]
+    dlParquetLastExportedTime: Optional[str]
+    sourceFirstCountedTime: Optional[str]
+    sourceLastCountedTime: Optional[str]
+
+
+class SubscriptionPatchRequestItem(SubscriptionStatePatchModel, total=False):
     name: Optional[str]
     plantStidCode: Optional[str]
-    plantSapCode: Optional[bool]
+    plantSapCode: Optional[str]
     terminal: Optional[str]
     fieldId: Optional[str]
     timeseriesId: Optional[str]
 
+
 class TimeseriesRequestItem(TypedDict, total=False):
     name: str
     description: Optional[str]
@@ -228,10 +248,12 @@ class StreamSubscriptionItemsModel(TypedDict):
 class StreamSubscriptionDataModel(TypedDict):
     data: StreamSubscriptionItemsModel
 
+
 class SubscriptionCounterModel(TypedDict):
     quota: int
     count: int
 
+
 class TimeseriesRequestFailedException(Exception):
     def __init__(self, response: Response) -> None:
         try:

tests/test_api.py

@@ -1,9 +1,13 @@
+import json
+
 import requests_mock
 import pytest
 from azure.core.credentials import AccessToken
 from omnia_timeseries import TimeseriesAPI, TimeseriesEnvironment
+from omnia_timeseries.api import IMSSubscriptionsManagementAPI
 from omnia_timeseries.http_client import TimeseriesRequestFailedException
 
+
 class DummyCredentials:
     def get_token(self, *scopes: str, **kwargs) -> AccessToken:
         return AccessToken("dummytoken", 0)
@@ -17,6 +21,17 @@ def api():
     return api
 
 
+@pytest.fixture
+def imssubscriptionsmanagementapi():
+    env = TimeseriesEnvironment.Dev()
+    api = IMSSubscriptionsManagementAPI(
+        azure_credential=DummyCredentials(),
+        environment=env,
+    )
+    api._base_url = "https://test"
+    return api
+
+
 def should_retry_request_when_failing_on_retryable_error_code_503_service_is_unavailable(
     api,
 ):
@@ -84,9 +99,9 @@ def should_forward_align_to_cache_parameter_for_multi_datapoints(api):
         api.get_multi_datapoints([{"id": "series-1"}], alignToCache=True)
         assert m.call_count == 1, "Expected a single request"
         request = m.request_history[0]
-        assert request.qs.get("aligntocache") == [
-            "true"
-        ], "alignToCache should be forwarded"
+        assert request.qs.get("aligntocache") == ["true"], (
+            "alignToCache should be forwarded"
+        )
 
 
 def should_forward_align_to_cache_parameter_for_get_aggregates(api):
@@ -99,9 +114,9 @@ def should_forward_align_to_cache_parameter_for_get_aggregates(api):
         api.get_aggregates("1234", ["avg"], alignToCache=False)
         assert m.call_count == 1, "Expected a single request"
         request = m.request_history[0]
-        assert request.qs.get("aligntocache") == [
-            "false"
-        ], "alignToCache should be forwarded as false"
+        assert request.qs.get("aligntocache") == ["false"], (
+            "alignToCache should be forwarded as false"
+        )
 
 
 def should_include_debug_query_param_when_debug_mode_is_enabled(api):
@@ -115,6 +130,101 @@ def should_include_debug_query_param_when_debug_mode_is_enabled(api):
         api.get_multi_datapoints([{"id": "series-1"}])
         assert m.call_count == 1, "Expected a single request"
         request = m.request_history[0]
-        assert request.qs.get("debug") == [
-            "true"
-        ], "Debug mode should inject debug=true"
+        assert request.qs.get("debug") == ["true"], (
+            "Debug mode should inject debug=true"
+        )
+
+
+def should_patch_subscription_by_uid_with_state_fields_only(
+    imssubscriptionsmanagementapi,
+):
+    payload = {
+        "tsdbEventsExported": 0,
+        "tsdbFirstExportedTime": None,
+        "tsdbLastExportedTime": None,
+        "dlEventsExported": 0,
+        "dlFirstExportedTime": None,
+        "dlLastExportedTime": None,
+        "sourceFirstCountedTime": None,
+        "sourceLastCountedTime": None,
+    }
+
+    with requests_mock.Mocker() as m:
+        m.register_uri(
+            "PATCH",
+            "https://test/uid/sub-123",
+            json={"data": {"items": []}, "count": 0},
+        )
+
+        response = imssubscriptionsmanagementapi.patch_subscription_by_uid(
+            uid="sub-123",
+            request=payload,
+        )
+
+        assert m.call_count == 1
+        request = m.request_history[0]
+        assert request.method == "PATCH"
+        assert json.loads(request.text) == payload
+        assert response["count"] == 0
+
+
+def should_patch_subscription_by_uid_with_subscription_and_state_fields(
+    imssubscriptionsmanagementapi,
+):
+    payload = {
+        "plantStidCode": "asdf-1234",
+        "plantSapCode": "1234",
+        "tsdbEventsExported": 0,
+        "tsdbFirstExportedTime": None,
+        "tsdbLastExportedTime": None,
+        "dlEventsExported": 0,
+        "dlFirstExportedTime": None,
+        "dlLastExportedTime": None,
+        "sourceFirstCountedTime": None,
+        "sourceLastCountedTime": None,
+    }
+
+    with requests_mock.Mocker() as m:
+        m.register_uri(
+            "PATCH",
+            "https://test/uid/sub-123",
+            json={"data": {"items": []}, "count": 0},
+        )
+
+        response = imssubscriptionsmanagementapi.patch_subscription_by_uid(
+            uid="sub-123",
+            request=payload,
+        )
+
+        assert m.call_count == 1
+        request = m.request_history[0]
+        assert request.method == "PATCH"
+        assert json.loads(request.text) == payload
+        assert response["count"] == 0
+
+
+def should_raise_on_patch_subscription_by_uid_when_api_returns_400(
+    imssubscriptionsmanagementapi,
+):
+    payload = {
+        "tsdbEventsExported": 1,
+        "notAllowedField": "boom",  # field intentionally unknown; API returns 400
+    }
+
+    with requests_mock.Mocker() as m:
+        m.register_uri(
+            "PATCH",
+            "https://test/uid/sub-123",
+            status_code=400,
+            text='{"message":"Invalid field: notAllowedField","traceId":"trace-1"}',
+        )
+
+        with pytest.raises(TimeseriesRequestFailedException) as exc:
+            imssubscriptionsmanagementapi.patch_subscription_by_uid(
+                uid="sub-123",
+                request=payload,
+            )
+
+        assert m.call_count == 1
+        assert exc.value.status_code == 400
+        assert "Invalid field" in exc.value.message