fix: complete finished_with_warnings and add cancelled to async RG status

HISTORY.rst

@@ -2,6 +2,12 @@
 History
 =======
 
+1.1.2 (2026-06-26)
+------------------
+
+* Completed wiring of the ``finished_with_warnings`` status for ``AsyncRulesetGenerationTaskStatus``.
+* Added the ``cancelled`` status to ``AsyncRulesetGenerationTaskStatus``.
+
 1.1.1 (2026-06-25)
 ------------------
 

datamasque/client/discovery.py

@@ -166,7 +166,7 @@ def get_generated_rulesets(self, connection_id: ConnectionId) -> list[Ruleset]:
             returns a list containing one ruleset
 
         Raises `AsyncRulesetGenerationInProgressError` if the task hasn't finished yet,
-        and `DataMasqueException` if it failed.
+        and `DataMasqueException` if it failed or was cancelled.
 
         Note that the ruleset(s) have autogenerated names, which you may want to customize before uploading.
         """
@@ -176,7 +176,11 @@ def get_generated_rulesets(self, connection_id: ConnectionId) -> list[Ruleset]:
             logger.error("Ruleset generation failed for connection: %s", connection_id)
             raise DataMasqueException(f"Ruleset generation failed for connection: {connection_id}")
 
-        if status is not AsyncRulesetGenerationTaskStatus.finished:
+        if status is AsyncRulesetGenerationTaskStatus.cancelled:
+            logger.error("Ruleset generation was cancelled for connection: %s", connection_id)
+            raise DataMasqueException(f"Ruleset generation was cancelled for connection: {connection_id}")
+
+        if not status.is_finished:
             logger.error(
                 "Ruleset generation is still in progress for connection: %s. Status: `%s`",
                 connection_id,

datamasque/client/models/status.py

@@ -64,15 +64,28 @@ class AsyncRulesetGenerationTaskStatus(enum.Enum):
     failed = "failed"
     running = "running"
     queued = "queued"
+    cancelled = "cancelled"
 
     @classmethod
     def get_final_states(cls) -> set["AsyncRulesetGenerationTaskStatus"]:
         """Returns the list of final statuses, i.e. the ruleset generation has completed, successfully or otherwise."""
 
-        return {cls.finished, cls.failed}
+        return {cls.finished, cls.finished_with_warnings, cls.failed, cls.cancelled}
+
+    @classmethod
+    def get_finished_states(cls) -> set["AsyncRulesetGenerationTaskStatus"]:
+        """Returns the list of statuses that indicate the ruleset generation completed successfully."""
+
+        return {cls.finished, cls.finished_with_warnings}
 
     @property
     def is_in_final_state(self) -> bool:
         """Returns True if this status is a final status."""
 
         return self in self.get_final_states()
+
+    @property
+    def is_finished(self) -> bool:
+        """Returns True if this status is a finished status."""
+
+        return self in self.get_finished_states()

tests/test_discovery.py

@@ -204,6 +204,47 @@ def test_get_generated_rulesets_success(client):
         assert rulesets[1].yaml == yaml_content_2.decode("utf-8")
 
 
+def test_get_generated_rulesets_finished_with_warnings_success(client):
+    """A task that finishes with warnings is still successful: its rulesets are returned, not treated as in-progress."""
+    connection_id = ConnectionId("1")
+    yaml_content = b"""
+    version: "1.0"
+    tasks:
+    - type: mask_table
+      table: table1
+      key: id
+      rules:
+      - column: col1
+        masks:
+        - type: do_nothing
+    """
+
+    with requests_mock.Mocker() as m:
+        m.get(
+            f"http://test-server/api/async-generate-ruleset/{connection_id}/",
+            json={"status": "finished_with_warnings"},
+            status_code=200,
+        )
+
+        zip_buffer = BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w") as zip_file:
+            zip_file.writestr("ruleset1.yml", yaml_content.decode("utf-8"))
+        zip_buffer.seek(0)
+
+        m.get(
+            f"http://test-server/api/async-generate-ruleset/{connection_id}/download-rulesets/",
+            content=zip_buffer.getvalue(),
+            headers={"Content-Disposition": 'attachment; filename="rulesets.zip"'},
+            status_code=200,
+        )
+
+        rulesets = client.get_generated_rulesets(connection_id)
+
+        assert len(rulesets) == 1
+        assert rulesets[0].name == "ruleset1"
+        assert rulesets[0].yaml == yaml_content.decode("utf-8")
+
+
 def test_get_generated_rulesets_empty_archive_raises(client):
     """A finished task whose download archive contains no ruleset files raises a clear error."""
     connection_id = ConnectionId("1")
@@ -290,6 +331,21 @@ def test_get_generated_rulesets_failed(client):
             client.get_generated_rulesets(connection_id)
 
 
+def test_get_generated_rulesets_cancelled(client):
+    """A cancelled task is terminal, so it raises a non-retryable DataMasqueException, not the in-progress error."""
+    connection_id = ConnectionId("1")
+
+    with requests_mock.Mocker() as m:
+        m.get(
+            f"http://test-server/api/async-generate-ruleset/{connection_id}/",
+            json={"status": "cancelled"},
+            status_code=200,
+        )
+
+        with pytest.raises(DataMasqueException, match="Ruleset generation was cancelled for connection"):
+            client.get_generated_rulesets(connection_id)
+
+
 def test_get_generated_rulesets_in_progress(client):
     connection_id = ConnectionId("1")