diff --git a/src/healthchecks_io/client/async_client.py b/src/healthchecks_io/client/async_client.py index 2076f9c..248c635 100644 --- a/src/healthchecks_io/client/async_client.py +++ b/src/healthchecks_io/client/async_client.py @@ -1,9 +1,11 @@ """An async healthchecks.io client.""" import asyncio +from types import TracebackType from typing import Dict from typing import List from typing import Optional from typing import Tuple +from typing import Type from httpx import AsyncClient as HTTPXAsyncClient @@ -56,6 +58,23 @@ class AsyncClient(AbstractClient): ] = f"py-healthchecks.io-async/{client_version}" self._client.headers["Content-type"] = "application/json" + async def __aenter__(self) -> "AsyncClient": + """Context manager entrance. + + Returns: + AsyncClient: returns this client as a context manager + """ + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + """Context manager exit.""" + await self._afinalizer_method() + def _finalizer_method(self) -> None: """Calls _afinalizer_method from a sync context to work with weakref.finalizer.""" asyncio.run(self._afinalizer_method()) diff --git a/src/healthchecks_io/client/sync_client.py b/src/healthchecks_io/client/sync_client.py index 3aa437e..4b4161c 100644 --- a/src/healthchecks_io/client/sync_client.py +++ b/src/healthchecks_io/client/sync_client.py @@ -1,8 +1,10 @@ """An async healthchecks.io client.""" +from types import TracebackType from typing import Dict from typing import List from typing import Optional from typing import Tuple +from typing import Type from httpx import Client as HTTPXClient @@ -50,6 +52,23 @@ class Client(AbstractClient): self._client.headers["user-agent"] = f"py-healthchecks.io/{client_version}" self._client.headers["Content-type"] = "application/json" + def __enter__(self) -> "Client": + """Context manager entrance. + + Returns: + Client: returns this client as a context manager + """ + return self + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: + """Context manager exit.""" + self._finalizer_method() + def _finalizer_method(self) -> None: """Closes the httpx client.""" self._client.close() diff --git a/tests/client/test_async.py b/tests/client/test_async.py index a9ecca7..6d79bf5 100644 --- a/tests/client/test_async.py +++ b/tests/client/test_async.py @@ -14,6 +14,42 @@ from healthchecks_io.client.exceptions import HCAPIAuthError from healthchecks_io.client.exceptions import HCAPIError +@pytest.mark.asyncio +@pytest.mark.respx +async def test_acreate_check_200_context_manager( + fake_check_api_result, respx_mock, test_async_client +): + checks_url = urljoin(test_async_client._api_url, "checks/") + respx_mock.post(checks_url).mock( + return_value=Response( + status_code=200, + json={ + "channels": "", + "desc": "", + "grace": 60, + "last_ping": None, + "n_pings": 0, + "name": "Backups", + "slug": "backups", + "next_ping": None, + "manual_resume": False, + "methods": "", + "pause_url": "https://healthchecks.io/api/v1/checks/f618072a-7bde-4eee-af63-71a77c5723bc/pause", + "ping_url": "https://hc-ping.com/f618072a-7bde-4eee-af63-71a77c5723bc", + "status": "new", + "tags": "prod www", + "timeout": 3600, + "update_url": "https://healthchecks.io/api/v1/checks/f618072a-7bde-4eee-af63-71a77c5723bc", + }, + ) + ) + async with test_async_client as test_client: + check = await test_client.create_check( + CheckCreate(name="test", tags="test", desc="test") + ) + assert check.name == "Backups" + + @pytest.mark.asyncio @pytest.mark.respx async def test_acreate_check_200(fake_check_api_result, respx_mock, test_async_client): diff --git a/tests/client/test_sync.py b/tests/client/test_sync.py index 853b656..a7585fb 100644 --- a/tests/client/test_sync.py +++ b/tests/client/test_sync.py @@ -14,6 +14,39 @@ from healthchecks_io.client.exceptions import HCAPIAuthError from healthchecks_io.client.exceptions import HCAPIError +@pytest.mark.respx +def test_create_check_200_context_manager( + fake_check_api_result, respx_mock, test_client +): + checks_url = urljoin(test_client._api_url, "checks/") + respx_mock.post(checks_url).mock( + return_value=Response( + status_code=200, + json={ + "channels": "", + "desc": "", + "grace": 60, + "last_ping": None, + "n_pings": 0, + "name": "Backups", + "slug": "backups", + "next_ping": None, + "manual_resume": False, + "methods": "", + "pause_url": "https://healthchecks.io/api/v1/checks/f618072a-7bde-4eee-af63-71a77c5723bc/pause", + "ping_url": "https://hc-ping.com/f618072a-7bde-4eee-af63-71a77c5723bc", + "status": "new", + "tags": "prod www", + "timeout": 3600, + "update_url": "https://healthchecks.io/api/v1/checks/f618072a-7bde-4eee-af63-71a77c5723bc", + }, + ) + ) + with test_client as tc: + check = tc.create_check(CheckCreate(name="test", tags="test", desc="test")) + assert check.name == "Backups" + + @pytest.mark.respx def test_create_check_200(fake_check_api_result, respx_mock, test_client): checks_url = urljoin(test_client._api_url, "checks/")