diff --git a/src/mixedbread/lib/parsing_jobs.py b/src/mixedbread/lib/parsing_jobs.py index 2985cf3..a5c087d 100644 --- a/src/mixedbread/lib/parsing_jobs.py +++ b/src/mixedbread/lib/parsing_jobs.py @@ -7,6 +7,7 @@ from . import polling from .._types import Omit, NotGiven, FileTypes, omit, not_given +from .._utils import is_given from .multipart_upload import MultipartUploadOptions from ..types.parsing.parsing_job import ParsingJob from ..types.parsing.element_type import ElementType @@ -35,13 +36,13 @@ def poll( **kwargs: Any, ) -> ParsingJob: """Poll a job's status until it reaches a terminal state.""" - polling_interval_ms = poll_interval_ms or 500 - polling_timeout_ms = poll_timeout_ms or None + polling_interval_ms = poll_interval_ms if is_given(poll_interval_ms) else 500 + polling_timeout_ms = poll_timeout_ms if is_given(poll_timeout_ms) else None return polling.poll( fn=functools.partial(self.retrieve, job_id, **kwargs), condition=lambda res: res.status in _TERMINAL, interval_seconds=polling_interval_ms / 1000, - timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms else None, + timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms is not None else None, ) def create_and_poll( @@ -123,13 +124,13 @@ async def poll( **kwargs: Any, ) -> ParsingJob: """Poll a job's status until it reaches a terminal state.""" - polling_interval_ms = poll_interval_ms or 500 - polling_timeout_ms = poll_timeout_ms or None + polling_interval_ms = poll_interval_ms if is_given(poll_interval_ms) else 500 + polling_timeout_ms = poll_timeout_ms if is_given(poll_timeout_ms) else None return await polling.poll_async( fn=functools.partial(self.retrieve, job_id, **kwargs), condition=lambda res: res.status in _TERMINAL, interval_seconds=polling_interval_ms / 1000, - timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms else None, + timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms is not None else None, ) async def create_and_poll( diff --git a/src/mixedbread/lib/polling.py b/src/mixedbread/lib/polling.py index f2aa95b..ac8f473 100644 --- a/src/mixedbread/lib/polling.py +++ b/src/mixedbread/lib/polling.py @@ -65,7 +65,7 @@ def poll( if max_attempts and attempt >= max_attempts: raise RuntimeError(f"Maximum attempts ({max_attempts}) reached") - if timeout_seconds: + if timeout_seconds is not None: elapsed = (datetime.now() - start_time).total_seconds() if elapsed >= timeout_seconds: raise TimeoutError(f"Timeout ({timeout_seconds}s) reached") @@ -114,7 +114,7 @@ async def poll_async( if max_attempts and attempt >= max_attempts: raise RuntimeError(f"Maximum attempts ({max_attempts}) reached") - if timeout_seconds: + if timeout_seconds is not None: elapsed = (datetime.now() - start_time).total_seconds() if elapsed >= timeout_seconds: raise TimeoutError(f"Timeout ({timeout_seconds}s) reached") diff --git a/src/mixedbread/lib/store_files.py b/src/mixedbread/lib/store_files.py index 682e9ed..3eedc41 100644 --- a/src/mixedbread/lib/store_files.py +++ b/src/mixedbread/lib/store_files.py @@ -12,6 +12,7 @@ from . import polling from .._types import Omit, NotGiven, FileTypes, omit, not_given +from .._utils import is_given from .multipart_upload import MultipartUploadOptions from ..types.stores.store_file import StoreFile from ..types.stores.store_file_config_param import StoreFileConfigParam @@ -48,13 +49,13 @@ def poll( Returns: The file object once it reaches a terminal state """ - polling_interval_ms = poll_interval_ms or 500 - polling_timeout_ms = poll_timeout_ms or None + polling_interval_ms = poll_interval_ms if is_given(poll_interval_ms) else 500 + polling_timeout_ms = poll_timeout_ms if is_given(poll_timeout_ms) else None return polling.poll( fn=functools.partial(self.retrieve, file_identifier, store_identifier=store_identifier, **kwargs), condition=lambda res: res.status in _TERMINAL, interval_seconds=polling_interval_ms / 1000, - timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms else None, + timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms is not None else None, ) def create_and_poll( @@ -163,13 +164,13 @@ async def poll( **kwargs: Any, ) -> StoreFile: """Poll a file's status until it reaches a terminal state.""" - polling_interval_ms = poll_interval_ms or 500 - polling_timeout_ms = poll_timeout_ms or None + polling_interval_ms = poll_interval_ms if is_given(poll_interval_ms) else 500 + polling_timeout_ms = poll_timeout_ms if is_given(poll_timeout_ms) else None return await polling.poll_async( fn=functools.partial(self.retrieve, file_identifier, store_identifier=store_identifier, **kwargs), condition=lambda res: res.status in _TERMINAL, interval_seconds=polling_interval_ms / 1000, - timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms else None, + timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms is not None else None, ) async def create_and_poll( diff --git a/src/mixedbread/lib/stores.py b/src/mixedbread/lib/stores.py index 6ee1e43..5e86c2f 100644 --- a/src/mixedbread/lib/stores.py +++ b/src/mixedbread/lib/stores.py @@ -12,6 +12,7 @@ from . import polling from .._types import Omit, NotGiven, omit, not_given +from .._utils import is_given from ..types.store import Store if TYPE_CHECKING: @@ -51,13 +52,13 @@ def poll( Returns: The store once it has settled """ - polling_interval_ms = poll_interval_ms or _DEFAULT_POLL_INTERVAL_MS - polling_timeout_ms = poll_timeout_ms or None + polling_interval_ms = poll_interval_ms if is_given(poll_interval_ms) else _DEFAULT_POLL_INTERVAL_MS + polling_timeout_ms = poll_timeout_ms if is_given(poll_timeout_ms) else None return polling.poll( fn=functools.partial(self.retrieve, store_identifier, **kwargs), condition=_is_settled, interval_seconds=polling_interval_ms / 1000, - timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms else None, + timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms is not None else None, ) def copy_and_poll( @@ -89,13 +90,13 @@ async def poll( **kwargs: Any, ) -> Store: """Poll a store until it is no longer ``in_progress``.""" - polling_interval_ms = poll_interval_ms or _DEFAULT_POLL_INTERVAL_MS - polling_timeout_ms = poll_timeout_ms or None + polling_interval_ms = poll_interval_ms if is_given(poll_interval_ms) else _DEFAULT_POLL_INTERVAL_MS + polling_timeout_ms = poll_timeout_ms if is_given(poll_timeout_ms) else None return await polling.poll_async( fn=functools.partial(self.retrieve, store_identifier, **kwargs), condition=_is_settled, interval_seconds=polling_interval_ms / 1000, - timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms else None, + timeout_seconds=polling_timeout_ms / 1000 if polling_timeout_ms is not None else None, ) async def copy_and_poll( diff --git a/tests/test_store_helpers.py b/tests/test_store_helpers.py index cecee29..8f24e69 100644 --- a/tests/test_store_helpers.py +++ b/tests/test_store_helpers.py @@ -1,5 +1,5 @@ from typing import List -from unittest.mock import Mock, AsyncMock +from unittest.mock import Mock, AsyncMock, patch import pytest @@ -46,6 +46,22 @@ def test_poll_returns_a_failed_copy() -> None: assert stores.poll("vs_copy", poll_interval_ms=1).status == "failed" +def test_poll_honors_zero_interval() -> None: + stores = _Stores(["in_progress", "completed"]) + + with patch("time.sleep") as sleep_mock: + assert stores.poll("vs_copy", poll_interval_ms=0).status == "completed" + + sleep_mock.assert_called_once_with(0.0) + + +def test_poll_honors_zero_timeout() -> None: + stores = _Stores(["in_progress", "in_progress", "completed"]) + + with pytest.raises(TimeoutError): + stores.poll("vs_copy", poll_interval_ms=1, poll_timeout_ms=0) + + @pytest.mark.asyncio async def test_async_copy_and_poll() -> None: stores = _AsyncStores(["in_progress", "completed"]) @@ -54,3 +70,21 @@ async def test_async_copy_and_poll() -> None: assert result.status == "completed" assert stores.retrieve_mock.await_count == 2 + + +@pytest.mark.asyncio +async def test_async_poll_honors_zero_interval() -> None: + stores = _AsyncStores(["in_progress", "completed"]) + + with patch("asyncio.sleep", new=AsyncMock()) as sleep_mock: + assert (await stores.poll("vs_copy", poll_interval_ms=0)).status == "completed" + + sleep_mock.assert_called_once_with(0.0) + + +@pytest.mark.asyncio +async def test_async_poll_honors_zero_timeout() -> None: + stores = _AsyncStores(["in_progress", "in_progress", "completed"]) + + with pytest.raises(TimeoutError): + await stores.poll("vs_copy", poll_interval_ms=1, poll_timeout_ms=0)