diff --git a/WrenchCL/Wrench/_notify.py b/WrenchCL/Wrench/_notify.py index 32e45f3..4846e51 100644 --- a/WrenchCL/Wrench/_notify.py +++ b/WrenchCL/Wrench/_notify.py @@ -130,6 +130,10 @@ def job_update( description: Optional[str] = None, *, workspace_id: Optional[str] = None, + status_note: Optional[str] = None, + input_rows: Optional[int] = None, + output_rows: Optional[int] = None, + ext_refs: Optional[dict] = None, base_url: Optional[str] = None, service_secret: Optional[str] = None, secret_env_var: str = "WRENCH_SERVICE_SECRET", @@ -148,6 +152,16 @@ def job_update( Progress percentage 0–100. description : str, optional Human-readable progress label shown alongside the percentage. + workspace_id : str, optional + The workspace (client) UUID. + status_note : str, optional + Status note to include in the update. + input_rows : int, optional + Number of input rows processed. + output_rows : int, optional + Number of output rows produced. + ext_refs : dict, optional + External references dictionary. Returns ------- @@ -170,6 +184,14 @@ def job_update( payload["progress_description"] = description if workspace_id is not None: payload["workspace_id"] = str(workspace_id) + if status_note is not None: + payload["status_note"] = status_note + if input_rows is not None: + payload["input_rows"] = input_rows + if output_rows is not None: + payload["output_rows"] = output_rows + if ext_refs is not None: + payload["ext_refs"] = ext_refs try: response = requests.patch( @@ -198,6 +220,10 @@ def job_close( source: str, *, notify: bool = True, + status_note: Optional[str] = None, + input_rows: Optional[int] = None, + output_rows: Optional[int] = None, + ext_refs: Optional[dict] = None, base_url: Optional[str] = None, service_secret: Optional[str] = None, secret_env_var: str = "WRENCH_SERVICE_SECRET", @@ -222,6 +248,14 @@ def job_close( Same source string used in job_register. notify : bool Whether to create a user-visible notification. Default True. + status_note : str, optional + Status note to include in the close event. + input_rows : int, optional + Number of input rows processed. + output_rows : int, optional + Number of output rows produced. + ext_refs : dict, optional + External references dictionary. Returns ------- @@ -247,6 +281,14 @@ def job_close( "source": source, "notify": notify, } + if status_note is not None: + payload["status_note"] = status_note + if input_rows is not None: + payload["input_rows"] = input_rows + if output_rows is not None: + payload["output_rows"] = output_rows + if ext_refs is not None: + payload["ext_refs"] = ext_refs try: response = requests.post( diff --git a/tests/test_wrench_notify.py b/tests/test_wrench_notify.py index edb5e32..f1953b5 100644 --- a/tests/test_wrench_notify.py +++ b/tests/test_wrench_notify.py @@ -324,6 +324,48 @@ def test_job_update_includes_workspace_id_when_provided(self, mock_patch): payload = call_kwargs["json"] assert payload["workspace_id"] == "ws-456" + @patch("WrenchCL.Wrench._notify.requests.patch") + def test_job_update_includes_all_lifecycle_fields(self, mock_patch): + mock_response = MagicMock() + mock_response.ok = True + mock_patch.return_value = mock_response + + job_update( + job_id="uuid-123", + progress=50, + input_rows=10, + output_rows=8, + status_note="done step", + ext_refs={"k": "v"}, + service_secret="test-secret" + ) + + call_kwargs = mock_patch.call_args[1] + payload = call_kwargs["json"] + assert payload["input_rows"] == 10 + assert payload["output_rows"] == 8 + assert payload["status_note"] == "done step" + assert payload["ext_refs"] == {"k": "v"} + + @patch("WrenchCL.Wrench._notify.requests.patch") + def test_job_update_omits_lifecycle_fields_when_not_provided(self, mock_patch): + mock_response = MagicMock() + mock_response.ok = True + mock_patch.return_value = mock_response + + job_update( + job_id="uuid-123", + progress=50, + service_secret="test-secret" + ) + + call_kwargs = mock_patch.call_args[1] + payload = call_kwargs["json"] + assert "input_rows" not in payload + assert "output_rows" not in payload + assert "status_note" not in payload + assert "ext_refs" not in payload + class TestJobClose: @patch("WrenchCL.Wrench._notify.requests.post") @@ -456,6 +498,32 @@ def test_job_close_default_notify_is_true(self, mock_post): payload = call_kwargs["json"] assert payload["notify"] is True + @patch("WrenchCL.Wrench._notify.requests.post") + def test_job_close_includes_all_lifecycle_fields(self, mock_post): + mock_response = MagicMock() + mock_response.ok = True + mock_post.return_value = mock_response + + job_close( + job_id="uuid-123", + workspace_id="ws-1", + status_code=200, + message="done", + source="elt", + input_rows=5, + output_rows=5, + status_note="ok", + ext_refs={"a": 1}, + service_secret="test-secret" + ) + + call_kwargs = mock_post.call_args[1] + payload = call_kwargs["json"] + assert payload["input_rows"] == 5 + assert payload["output_rows"] == 5 + assert payload["status_note"] == "ok" + assert payload["ext_refs"] == {"a": 1} + class TestAutoArnIntegration: @patch("WrenchCL.Wrench._notify.requests.post")