magic_hour 0.34.0__py3-none-any.whl → 0.36.0__py3-none-any.whl

magic_hour/resources/v1/image_projects/client.py

@@ -1,3 +1,8 @@
+import asyncio
+import logging
+import os
+import pydantic
+import time
 import typing
 
 from magic_hour.core import (
@@ -6,13 +11,87 @@ from magic_hour.core import (
     SyncBaseClient,
     default_request_options,
 )
+from magic_hour.helpers.download import download_files_async, download_files_sync
 from magic_hour.types import models
 
 
+logger = logging.getLogger(__name__)
+
+
+class V1ImageProjectsGetResponseWithDownloads(models.V1ImageProjectsGetResponse):
+    downloaded_paths: typing.Optional[typing.List[str]] = pydantic.Field(
+        default=None, alias="downloaded_paths"
+    )
+    """
+    The paths to the downloaded files.
+
+    This field is only populated if `download_outputs` is True and the image project is complete.
+    """
+
+
 class ImageProjectsClient:
     def __init__(self, *, base_client: SyncBaseClient):
         self._base_client = base_client
 
+    def check_result(
+        self,
+        id: str,
+        wait_for_completion: bool,
+        download_outputs: bool,
+        download_directory: typing.Optional[str] = None,
+    ) -> V1ImageProjectsGetResponseWithDownloads:
+        """
+        Check the result of an image project with optional waiting and downloading.
+
+        This method retrieves the status of an image project and optionally waits for completion
+        and downloads the output files.
+
+        Args:
+            id: Unique ID of the image project
+            wait_for_completion: Whether to wait for the image project to complete
+            download_outputs: Whether to download the outputs
+            download_directory: The directory to download the outputs to. If not provided,
+                the outputs will be downloaded to the current working directory
+
+        Returns:
+            V1ImageProjectsGetResponseWithDownloads: The image project response with optional
+                downloaded file paths included
+        """
+        api_response = self.get(id=id)
+        if not wait_for_completion:
+            response = V1ImageProjectsGetResponseWithDownloads(
+                **api_response.model_dump()
+            )
+            return response
+
+        poll_interval = float(os.getenv("MAGIC_HOUR_POLL_INTERVAL", "0.5"))
+
+        status = api_response.status
+
+        while status not in ["complete", "error", "canceled"]:
+            api_response = self.get(id=id)
+            status = api_response.status
+            time.sleep(poll_interval)
+
+        if api_response.status != "complete":
+            log = logger.error if api_response.status == "error" else logger.info
+            log(
+                f"Image project {id} has status {api_response.status}: {api_response.error}"
+            )
+            return V1ImageProjectsGetResponseWithDownloads(**api_response.model_dump())
+
+        if not download_outputs:
+            return V1ImageProjectsGetResponseWithDownloads(**api_response.model_dump())
+
+        downloaded_paths = download_files_sync(
+            downloads=api_response.downloads,
+            download_directory=download_directory,
+        )
+
+        return V1ImageProjectsGetResponseWithDownloads(
+            **api_response.model_dump(), downloaded_paths=downloaded_paths
+        )
+
     def delete(
         self, *, id: str, request_options: typing.Optional[RequestOptions] = None
     ) -> None:
@@ -95,6 +174,65 @@ class AsyncImageProjectsClient:
     def __init__(self, *, base_client: AsyncBaseClient):
         self._base_client = base_client
 
+    async def check_result(
+        self,
+        id: str,
+        wait_for_completion: bool,
+        download_outputs: bool,
+        download_directory: typing.Optional[str] = None,
+    ) -> V1ImageProjectsGetResponseWithDownloads:
+        """
+        Check the result of an image project with optional waiting and downloading.
+
+        This method retrieves the status of an image project and optionally waits for completion
+        and downloads the output files.
+
+        Args:
+            id: Unique ID of the image project
+            wait_for_completion: Whether to wait for the image project to complete
+            download_outputs: Whether to download the outputs
+            download_directory: The directory to download the outputs to. If not provided,
+                the outputs will be downloaded to the current working directory
+
+        Returns:
+            V1ImageProjectsGetResponseWithDownloads: The image project response with optional
+                downloaded file paths included
+        """
+        api_response = await self.get(id=id)
+        if not wait_for_completion:
+            response = V1ImageProjectsGetResponseWithDownloads(
+                **api_response.model_dump()
+            )
+            return response
+
+        poll_interval = float(os.getenv("MAGIC_HOUR_POLL_INTERVAL", "0.5"))
+
+        status = api_response.status
+
+        while status not in ["complete", "error", "canceled"]:
+            api_response = await self.get(id=id)
+            status = api_response.status
+            await asyncio.sleep(poll_interval)
+
+        if api_response.status != "complete":
+            log = logger.error if api_response.status == "error" else logger.info
+            log(
+                f"Image project {id} has status {api_response.status}: {api_response.error}"
+            )
+            return V1ImageProjectsGetResponseWithDownloads(**api_response.model_dump())
+
+        if not download_outputs:
+            return V1ImageProjectsGetResponseWithDownloads(**api_response.model_dump())
+
+        downloaded_paths = await download_files_async(
+            downloads=api_response.downloads,
+            download_directory=download_directory,
+        )
+
+        return V1ImageProjectsGetResponseWithDownloads(
+            **api_response.model_dump(), downloaded_paths=downloaded_paths
+        )
+
     async def delete(
         self, *, id: str, request_options: typing.Optional[RequestOptions] = None
     ) -> None:

magic_hour/resources/v1/image_projects/client_test.py

@@ -0,0 +1,527 @@
+import datetime
+import pytest
+import httpx
+from pathlib import Path
+from typing import Any, Generator, Literal, Union, List
+from unittest.mock import Mock, AsyncMock
+
+from magic_hour.types import models
+from magic_hour.resources.v1.image_projects.client import (
+    ImageProjectsClient,
+    AsyncImageProjectsClient,
+)
+
+
+class DummyResponse(models.V1ImageProjectsGetResponse):
+    """Helper response with defaults"""
+
+    def __init__(
+        self,
+        *,
+        status: Literal[
+            "complete", "queued", "rendering", "error", "canceled"
+        ] = "complete",
+        download_url: Union[str, None] = None,
+        error: Union[str, None] = None,
+    ):
+        # Create error object if error string is provided
+        error_obj = None
+        if error:
+            error_obj = models.V1ImageProjectsGetResponseError(
+                code="TEST_ERROR", message=error
+            )
+
+        super().__init__(
+            id="test-id",
+            created_at=datetime.datetime.now().isoformat(),
+            credits_charged=0,
+            enabled=True,
+            image_count=0,
+            name="test-name",
+            total_frame_cost=0,
+            type="test-type",
+            status=status,
+            downloads=[
+                models.V1ImageProjectsGetResponseDownloadsItem(
+                    url=download_url, expires_at="2024-01-01T00:00:00Z"
+                )
+            ]
+            if download_url
+            else [],
+            error=error_obj,
+        )
+
+
+@pytest.fixture
+def mock_base_client() -> Generator[Mock, None, None]:
+    yield Mock()
+
+
+@pytest.fixture
+def mock_async_base_client() -> Generator[AsyncMock, None, None]:
+    yield AsyncMock()
+
+
+def test_delete_calls_base_client(mock_base_client: Mock) -> None:
+    client = ImageProjectsClient(base_client=mock_base_client)
+    client.delete(id="123")
+
+    mock_base_client.request.assert_called_once()
+    call = mock_base_client.request.call_args[1]
+    assert call["method"] == "DELETE"
+    assert "/v1/image-projects/123" in call["path"]
+
+
+def test_get_calls_base_client(mock_base_client: Mock) -> None:
+    client = ImageProjectsClient(base_client=mock_base_client)
+    mock_base_client.request.return_value = DummyResponse()
+
+    resp = client.get(id="abc")
+
+    mock_base_client.request.assert_called_once()
+    assert isinstance(resp, models.V1ImageProjectsGetResponse)
+    assert resp.id == "test-id"
+
+
+def test_check_result_no_wait_no_download(mock_base_client: Mock) -> None:
+    client = ImageProjectsClient(base_client=mock_base_client)
+    mock_base_client.request.return_value = DummyResponse(status="queued")
+
+    resp = client.check_result(
+        id="xyz",
+        wait_for_completion=False,
+        download_outputs=False,
+    )
+
+    assert resp.downloaded_paths is None
+
+
+def test_check_result_wait_until_complete(
+    monkeypatch: Any, mock_base_client: Mock
+) -> None:
+    client = ImageProjectsClient(base_client=mock_base_client)
+
+    # First calls return queued, then complete
+    mock_base_client.request.side_effect = [
+        DummyResponse(status="queued"),
+        DummyResponse(status="queued"),
+        DummyResponse(status="complete"),
+    ]
+
+    monkeypatch.setattr("time.sleep", lambda _: None)  # type: ignore
+
+    resp = client.check_result(
+        id="xyz", wait_for_completion=True, download_outputs=False
+    )
+
+    assert resp.status == "complete"
+    assert resp.downloaded_paths is None
+
+
+def test_check_result_download_outputs(
+    tmp_path: Path, mock_base_client: Mock, monkeypatch: Any
+) -> None:
+    client = ImageProjectsClient(base_client=mock_base_client)
+
+    file_url = "https://example.com/file.png"
+    mock_base_client.request.return_value = DummyResponse(
+        status="complete",
+        download_url=file_url,
+    )
+
+    # Create a mock response for httpx
+    mock_request = httpx.Request("GET", "https://example.com/file.png")
+    mock_response = httpx.Response(200, content=b"fake png", request=mock_request)
+
+    # Mock the httpx.Client class
+    class MockClient:
+        def __init__(self):
+            pass
+
+        def __enter__(self) -> "MockClient":
+            return self
+
+        def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+            pass
+
+        def get(self, url: str) -> httpx.Response:
+            return mock_response
+
+    monkeypatch.setattr(httpx, "Client", MockClient)
+
+    resp = client.check_result(
+        id="xyz",
+        wait_for_completion=True,
+        download_outputs=True,
+        download_directory=str(tmp_path),
+    )
+
+    assert resp.status == "complete"
+    assert resp.downloaded_paths
+    saved_file = Path(resp.downloaded_paths[0])
+    assert saved_file.exists()
+    assert saved_file.read_bytes() == b"fake png"
+
+
+def test_check_result_error_status(mock_base_client: Mock) -> None:
+    client = ImageProjectsClient(base_client=mock_base_client)
+    mock_base_client.request.return_value = DummyResponse(status="error", error="Boom!")
+
+    resp = client.check_result(
+        id="err", wait_for_completion=True, download_outputs=False
+    )
+    assert resp.status == "error"
+    assert resp.error is not None
+    assert resp.error.message == "Boom!"
+    assert resp.downloaded_paths is None
+
+
+def test_check_result_canceled_status(mock_base_client: Mock) -> None:
+    client = ImageProjectsClient(base_client=mock_base_client)
+    mock_base_client.request.return_value = DummyResponse(status="canceled")
+
+    resp = client.check_result(
+        id="cancel", wait_for_completion=True, download_outputs=False
+    )
+    assert resp.status == "canceled"
+    assert resp.downloaded_paths is None
+
+
+def test_check_result_poll_interval_default(
+    mock_base_client: Mock, monkeypatch: Any
+) -> None:
+    client = ImageProjectsClient(base_client=mock_base_client)
+
+    # First calls return queued, then complete
+    mock_base_client.request.side_effect = [
+        DummyResponse(status="queued"),
+        DummyResponse(status="complete"),
+    ]
+
+    # Mock time.sleep to track calls
+    sleep_calls: List[float] = []
+
+    def mock_sleep(seconds: float) -> None:
+        sleep_calls.append(seconds)
+
+    monkeypatch.setattr("time.sleep", mock_sleep)
+
+    resp = client.check_result(
+        id="xyz", wait_for_completion=True, download_outputs=False
+    )
+
+    assert resp.status == "complete"
+    # Should have slept once with default interval (0.5)
+    assert len(sleep_calls) == 1
+    assert sleep_calls[0] == 0.5
+
+
+def test_check_result_poll_interval_custom(
+    mock_base_client: Mock, monkeypatch: Any
+) -> None:
+    client = ImageProjectsClient(base_client=mock_base_client)
+
+    # Set custom poll interval
+    monkeypatch.setenv("MAGIC_HOUR_POLL_INTERVAL", "1.0")
+
+    # First calls return queued, then complete
+    mock_base_client.request.side_effect = [
+        DummyResponse(status="queued"),
+        DummyResponse(status="complete"),
+    ]
+
+    # Mock time.sleep to track calls
+    sleep_calls: List[float] = []
+
+    def mock_sleep(seconds: float) -> None:
+        sleep_calls.append(seconds)
+
+    monkeypatch.setattr("time.sleep", mock_sleep)
+
+    resp = client.check_result(
+        id="xyz", wait_for_completion=True, download_outputs=False
+    )
+
+    assert resp.status == "complete"
+    # Should have slept once with custom interval (1.0)
+    assert len(sleep_calls) == 1
+    assert sleep_calls[0] == 1.0
+
+
+def test_check_result_poll_interval_multiple_polls(
+    mock_base_client: Mock, monkeypatch: Any
+) -> None:
+    client = ImageProjectsClient(base_client=mock_base_client)
+
+    # Set custom poll interval
+    monkeypatch.setenv("MAGIC_HOUR_POLL_INTERVAL", "0.1")
+
+    # Multiple calls return queued before complete
+    mock_base_client.request.side_effect = [
+        DummyResponse(status="queued"),
+        DummyResponse(status="queued"),
+        DummyResponse(status="queued"),
+        DummyResponse(status="complete"),
+    ]
+
+    # Mock time.sleep to track calls
+    sleep_calls: List[float] = []
+
+    def mock_sleep(seconds: float) -> None:
+        sleep_calls.append(seconds)
+
+    monkeypatch.setattr("time.sleep", mock_sleep)
+
+    resp = client.check_result(
+        id="xyz", wait_for_completion=True, download_outputs=False
+    )
+
+    assert resp.status == "complete"
+    # Should have slept 3 times with custom interval (0.1)
+    assert len(sleep_calls) == 3
+    assert all(sleep_time == 0.1 for sleep_time in sleep_calls)
+
+
+@pytest.mark.asyncio
+async def test_async_delete_calls_base_client(
+    mock_async_base_client: AsyncMock,
+) -> None:
+    client = AsyncImageProjectsClient(base_client=mock_async_base_client)
+    await client.delete(id="456")
+
+    mock_async_base_client.request.assert_called_once()
+    call = mock_async_base_client.request.call_args[1]
+    assert call["method"] == "DELETE"
+    assert "/v1/image-projects/456" in call["path"]
+
+
+@pytest.mark.asyncio
+async def test_async_get_calls_base_client(mock_async_base_client: AsyncMock) -> None:
+    client = AsyncImageProjectsClient(base_client=mock_async_base_client)
+    mock_async_base_client.request.return_value = DummyResponse()
+
+    resp = await client.get(id="zzz")
+
+    mock_async_base_client.request.assert_called_once()
+    assert isinstance(resp, models.V1ImageProjectsGetResponse)
+    assert resp.id == "test-id"
+
+
+@pytest.mark.asyncio
+async def test_async_check_result_no_wait_no_download(
+    mock_async_base_client: AsyncMock,
+) -> None:
+    client = AsyncImageProjectsClient(base_client=mock_async_base_client)
+    mock_async_base_client.request.return_value = DummyResponse(status="queued")
+
+    resp = await client.check_result(
+        id="xyz",
+        wait_for_completion=False,
+        download_outputs=False,
+    )
+
+    assert resp.downloaded_paths is None
+
+
+@pytest.mark.asyncio
+async def test_async_check_result_wait_until_complete(
+    mock_async_base_client: AsyncMock, monkeypatch: Any
+) -> None:
+    client = AsyncImageProjectsClient(base_client=mock_async_base_client)
+
+    # First calls return queued, then complete
+    mock_async_base_client.request.side_effect = [
+        DummyResponse(status="queued"),
+        DummyResponse(status="queued"),
+        DummyResponse(status="complete"),
+    ]
+
+    sleep_calls: List[float] = []
+
+    async def async_mock_sleep(seconds: float) -> None:
+        sleep_calls.append(seconds)
+
+    monkeypatch.setattr("asyncio.sleep", async_mock_sleep)
+
+    resp = await client.check_result(
+        id="xyz", wait_for_completion=True, download_outputs=False
+    )
+
+    assert resp.status == "complete"
+    assert resp.downloaded_paths is None
+
+
+@pytest.mark.asyncio
+async def test_async_check_result_download_outputs(
+    tmp_path: Path, mock_async_base_client: AsyncMock, monkeypatch: Any
+) -> None:
+    client = AsyncImageProjectsClient(base_client=mock_async_base_client)
+
+    file_url = "https://example.com/file.png"
+    mock_async_base_client.request.return_value = DummyResponse(
+        status="complete",
+        download_url=file_url,
+    )
+
+    # Create a mock response for httpx
+    mock_request = httpx.Request("GET", "https://example.com/file.png")
+    mock_response = httpx.Response(200, content=b"fake png", request=mock_request)
+
+    # Mock the httpx.AsyncClient class
+    class MockAsyncClient:
+        def __init__(self):
+            pass
+
+        async def __aenter__(self) -> "MockAsyncClient":
+            return self
+
+        async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+            pass
+
+        async def get(self, url: str) -> httpx.Response:
+            return mock_response
+
+    monkeypatch.setattr(httpx, "AsyncClient", MockAsyncClient)
+
+    resp = await client.check_result(
+        id="xyz",
+        wait_for_completion=True,
+        download_outputs=True,
+        download_directory=str(tmp_path),
+    )
+
+    assert resp.status == "complete"
+    assert resp.downloaded_paths
+    saved_file = Path(resp.downloaded_paths[0])
+    assert saved_file.exists()
+    assert saved_file.read_bytes() == b"fake png"
+
+
+@pytest.mark.asyncio
+async def test_async_check_result_error_status(
+    mock_async_base_client: AsyncMock,
+) -> None:
+    client = AsyncImageProjectsClient(base_client=mock_async_base_client)
+    mock_async_base_client.request.return_value = DummyResponse(
+        status="error", error="Boom!"
+    )
+
+    resp = await client.check_result(
+        id="err", wait_for_completion=True, download_outputs=False
+    )
+    assert resp.status == "error"
+    assert resp.error is not None
+    assert resp.error.message == "Boom!"
+    assert resp.downloaded_paths is None
+
+
+@pytest.mark.asyncio
+async def test_async_check_result_canceled_status(
+    mock_async_base_client: AsyncMock,
+) -> None:
+    client = AsyncImageProjectsClient(base_client=mock_async_base_client)
+    mock_async_base_client.request.return_value = DummyResponse(status="canceled")
+
+    resp = await client.check_result(
+        id="cancel", wait_for_completion=True, download_outputs=False
+    )
+    assert resp.status == "canceled"
+    assert resp.downloaded_paths is None
+
+
+@pytest.mark.asyncio
+async def test_async_check_result_poll_interval_default(
+    mock_async_base_client: AsyncMock, monkeypatch: Any
+) -> None:
+    client = AsyncImageProjectsClient(base_client=mock_async_base_client)
+
+    # First calls return queued, then complete
+    mock_async_base_client.request.side_effect = [
+        DummyResponse(status="queued"),
+        DummyResponse(status="complete"),
+    ]
+
+    # Mock time.sleep to track calls
+    sleep_calls: List[float] = []
+
+    async def async_mock_sleep(seconds: float) -> None:
+        sleep_calls.append(seconds)
+
+    monkeypatch.setattr("asyncio.sleep", async_mock_sleep)
+
+    resp = await client.check_result(
+        id="xyz", wait_for_completion=True, download_outputs=False
+    )
+
+    assert resp.status == "complete"
+    # Should have slept once with default interval (0.5)
+    assert len(sleep_calls) == 1
+    assert sleep_calls[0] == 0.5
+
+
+@pytest.mark.asyncio
+async def test_async_check_result_poll_interval_custom(
+    mock_async_base_client: AsyncMock, monkeypatch: Any
+) -> None:
+    client = AsyncImageProjectsClient(base_client=mock_async_base_client)
+
+    # Set custom poll interval
+    monkeypatch.setenv("MAGIC_HOUR_POLL_INTERVAL", "2.0")
+
+    # First calls return queued, then complete
+    mock_async_base_client.request.side_effect = [
+        DummyResponse(status="queued"),
+        DummyResponse(status="complete"),
+    ]
+
+    # Mock time.sleep to track calls
+    sleep_calls: List[float] = []
+
+    async def async_mock_sleep(seconds: float) -> None:
+        sleep_calls.append(seconds)
+
+    monkeypatch.setattr("asyncio.sleep", async_mock_sleep)
+
+    resp = await client.check_result(
+        id="xyz", wait_for_completion=True, download_outputs=False
+    )
+
+    assert resp.status == "complete"
+    # Should have slept once with custom interval (2.0)
+    assert len(sleep_calls) == 1
+    assert sleep_calls[0] == 2.0
+
+
+@pytest.mark.asyncio
+async def test_async_check_result_poll_interval_multiple_polls(
+    mock_async_base_client: AsyncMock, monkeypatch: Any
+) -> None:
+    client = AsyncImageProjectsClient(base_client=mock_async_base_client)
+
+    # Set custom poll interval
+    monkeypatch.setenv("MAGIC_HOUR_POLL_INTERVAL", "0.3")
+
+    # Multiple calls return queued before complete
+    mock_async_base_client.request.side_effect = [
+        DummyResponse(status="queued"),
+        DummyResponse(status="queued"),
+        DummyResponse(status="queued"),
+        DummyResponse(status="complete"),
+    ]
+
+    # Mock time.sleep to track calls
+    sleep_calls: List[float] = []
+
+    async def async_mock_sleep(seconds: float) -> None:
+        sleep_calls.append(seconds)
+
+    monkeypatch.setattr("asyncio.sleep", async_mock_sleep)
+
+    resp = await client.check_result(
+        id="xyz", wait_for_completion=True, download_outputs=False
+    )
+
+    assert resp.status == "complete"
+    # Should have slept 3 times with custom interval (0.3)
+    assert len(sleep_calls) == 3
+    assert all(sleep_time == 0.3 for sleep_time in sleep_calls)