diff --git a/getstream/stream.py b/getstream/stream.py index 0d080c6f..ec946b01 100644 --- a/getstream/stream.py +++ b/getstream/stream.py @@ -1,5 +1,6 @@ from __future__ import annotations +from contextlib import AsyncExitStack from functools import cached_property import time from typing import List, Optional @@ -207,6 +208,20 @@ def moderation(self) -> AsyncModerationClient: user_agent=self.user_agent, ) + async def aclose(self): + """Close all child clients and the main HTTPX client.""" + # AsyncExitStack ensures all clients are closed even if one fails. + # video/chat/moderation are @cached_property - only close if accessed. + async with AsyncExitStack() as stack: + cached = self.__dict__ + if "video" in cached: + stack.push_async_callback(self.video.aclose) + if "chat" in cached: + stack.push_async_callback(self.chat.aclose) + if "moderation" in cached: + stack.push_async_callback(self.moderation.aclose) + stack.push_async_callback(super().aclose) + @cached_property def feeds(self): raise NotImplementedError("Feeds not supported for async client") diff --git a/tests/test_async_stream_close.py b/tests/test_async_stream_close.py new file mode 100644 index 00000000..c9445185 --- /dev/null +++ b/tests/test_async_stream_close.py @@ -0,0 +1,43 @@ +import pytest + +from getstream import AsyncStream + + +@pytest.mark.asyncio +class TestAsyncStreamClose: + async def test_aclose_closes_main_client(self): + client = AsyncStream(api_key="fake", api_secret="fake") + + assert client.client.is_closed is False + await client.aclose() + assert client.client.is_closed is True + + async def test_aclose_closes_video_client(self): + client = AsyncStream(api_key="fake", api_secret="fake") + _ = client.video # trigger cached_property + + assert client.video.client.is_closed is False + await client.aclose() + assert client.video.client.is_closed is True + + async def test_aclose_closes_chat_client(self): + client = AsyncStream(api_key="fake", api_secret="fake") + _ = client.chat + + assert client.chat.client.is_closed is False + await client.aclose() + assert client.chat.client.is_closed is True + + async def test_aclose_closes_moderation_client(self): + client = AsyncStream(api_key="fake", api_secret="fake") + _ = client.moderation + + assert client.moderation.client.is_closed is False + await client.aclose() + assert client.moderation.client.is_closed is True + + async def test_aclose_without_child_clients(self): + """aclose() should work even if video/chat were never accessed.""" + client = AsyncStream(api_key="fake", api_secret="fake") + await client.aclose() + assert client.client.is_closed is True