From 66e9cfdd16059cbabed2d9a19117b3e27c999f37 Mon Sep 17 00:00:00 2001 From: YBubu Date: Mon, 25 Nov 2024 17:12:44 +0000 Subject: [PATCH 1/3] Fix ErrorResponse validation error and handle reponse.success = False --- pyproject.toml | 2 +- workflowai/core/client/api.py | 8 ++- workflowai/core/client/api_test.py | 85 ++++++++++++++++++++++++++++++ 3 files changed, 92 insertions(+), 3 deletions(-) create mode 100644 workflowai/core/client/api_test.py diff --git a/pyproject.toml b/pyproject.toml index bb79da5..4dd7477 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "workflowai" -version = "0.4.1" +version = "0.4.2" description = "" authors = ["Guillaume Aquilina "] readme = "README.md" diff --git a/workflowai/core/client/api.py b/workflowai/core/client/api.py index 5c53fd5..5416387 100644 --- a/workflowai/core/client/api.py +++ b/workflowai/core/client/api.py @@ -1,4 +1,3 @@ -from json import JSONDecodeError from typing import Any, AsyncIterator, Literal, Optional, TypeVar, Union, overload import httpx @@ -99,7 +98,7 @@ def _extract_error( try: res = ErrorResponse.model_validate_json(data) return WorkflowAIError(error=res.error, task_run_id=res.task_run_id, response=response) - except JSONDecodeError: + except ValidationError: raise WorkflowAIError( error=BaseError( message="Unknown error" if exception is None else str(exception), @@ -123,6 +122,11 @@ async def stream( content=data.model_dump_json(exclude_none=True), headers={"Content-Type": "application/json"}, ) as response: + if not response.is_success: + # We need to read the response to get the error message + await response.aread() + response.raise_for_status() + async for chunk in response.aiter_bytes(): payload = "" try: diff --git a/workflowai/core/client/api_test.py b/workflowai/core/client/api_test.py new file mode 100644 index 0000000..8f0f6db --- /dev/null +++ b/workflowai/core/client/api_test.py @@ -0,0 +1,85 @@ +import httpx +import pytest +from pydantic import BaseModel +from pytest_httpx import HTTPXMock + +from workflowai.core.client.api import APIClient +from workflowai.core.domain.errors import WorkflowAIError + + +class TestAPIClientExtractError: + def test_extract_error(self): + client = APIClient(endpoint="test_endpoint", api_key="test_api_key") + + # Test valid JSON error response + response = httpx.Response( + status_code=400, + json={ + "error": { + "message": "Test error message", + "details": {"key": "value"}, + }, + "task_run_id": "test_task_123", + }, + ) + + error = client._extract_error(response, response.content) # pyright:ignore[reportPrivateUsage] + assert isinstance(error, WorkflowAIError) + assert error.error.message == "Test error message" + assert error.error.details == {"key": "value"} + assert error.task_run_id == "test_task_123" + assert error.response == response + + def test_extract_error_invalid_json(self): + client = APIClient(endpoint="test_endpoint", api_key="test_api_key") + + # Test invalid JSON response + invalid_data = b"Invalid JSON data" + response = httpx.Response(status_code=400, content=invalid_data) + + with pytest.raises(WorkflowAIError) as e: + client._extract_error(response, invalid_data) # pyright:ignore[reportPrivateUsage] + assert isinstance(e.value, WorkflowAIError) + assert e.value.error.message == "Unknown error" + assert e.value.error.details == {"raw": "b'Invalid JSON data'"} + assert e.value.response == response + + def test_extract_error_with_custom_error(self): + client = APIClient(endpoint="test_endpoint", api_key="test_api_key") + + # Test with provided exception + invalid_data = "{'detail': 'Not Found'}" + response = httpx.Response(status_code=404, content=invalid_data) + exception = ValueError("Custom error") + + with pytest.raises(WorkflowAIError) as e: + client._extract_error(response, invalid_data, exception) # pyright:ignore[reportPrivateUsage] + assert isinstance(e.value, WorkflowAIError) + assert e.value.error.message == "Custom error" + assert e.value.error.details == {"raw": "{'detail': 'Not Found'}"} + assert e.value.response == response + + +async def test_stream_404(httpx_mock: HTTPXMock): + class TestInputModel(BaseModel): + test_input: str + + class TestOutputModel(BaseModel): + test_output: str + + httpx_mock.add_response(status_code=404) + + client = APIClient(endpoint="https://blabla.com", api_key="test_api_key") + + try: + async for _ in client.stream( + method="GET", + path="test_path", + data=TestInputModel(test_input="test"), + returns=TestOutputModel, + ): + pass + except Exception as e: + assert isinstance(e, httpx.HTTPStatusError) + assert e.response.status_code == 404 + assert e.response.reason_phrase == "Not Found" From 08645b86de45024b30c155e638a9638f89edfaf4 Mon Sep 17 00:00:00 2001 From: YBubu Date: Mon, 25 Nov 2024 17:15:20 +0000 Subject: [PATCH 2/3] change ruff rules --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 4dd7477..a3d9039 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,6 +51,8 @@ ignore = [ "TD", "PYI051", "FIX002", + "SLF001", #reportPrivateUsage + "PT017", # Do not force using pytest.raises ] # Allow fix for all enabled rules (when `--fix`) is provided. From 36a153ca0b1436e3d1df6c87372f8b03e1097ca7 Mon Sep 17 00:00:00 2001 From: YBubu Date: Mon, 25 Nov 2024 17:16:36 +0000 Subject: [PATCH 3/3] llint --- workflowai/core/client/api_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workflowai/core/client/api_test.py b/workflowai/core/client/api_test.py index 8f0f6db..b346369 100644 --- a/workflowai/core/client/api_test.py +++ b/workflowai/core/client/api_test.py @@ -79,7 +79,7 @@ class TestOutputModel(BaseModel): returns=TestOutputModel, ): pass - except Exception as e: + except httpx.HTTPStatusError as e: assert isinstance(e, httpx.HTTPStatusError) assert e.response.status_code == 404 assert e.response.reason_phrase == "Not Found"