"""Tests for LLM provider wrapper classes."""

from __future__ import annotations

from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

from tokpulse.tracker import TokpulseTracker
from tokpulse.wrappers import (
    AnthropicBetaWrapper,
    AnthropicMessagesWrapper,
    AnthropicWrapper,
    AsyncAnthropicBetaWrapper,
    AsyncAnthropicMessagesWrapper,
    AsyncAnthropicWrapper,
    AsyncOpenAICompletionsWrapper,
    AsyncOpenAIWrapper,
    GoogleWrapper,
    OpenAICompletionsWrapper,
)


def _mock_tracker() -> TokpulseTracker:
    """Return a tracker with background threads stopped to avoid network calls."""
    tracker = TokpulseTracker(
        api_key="lct_test",
        api_url="http://localhost:9999",
        allow_insecure_http=True,
        flush_interval_seconds=3600,
    )
    # Stop the background threads so they do not interfere with unit tests.
    tracker._stop_event.set()
    tracker._flush_trigger.set()
    return tracker


class TestAnthropicMessagesWrapper:
    def _make_wrapper(
        self, tracker: TokpulseTracker | None = None
    ) -> tuple[AnthropicMessagesWrapper, TokpulseTracker]:
        if tracker is None:
            tracker = _mock_tracker()
        mock_messages = MagicMock()
        return AnthropicMessagesWrapper(mock_messages, tracker), tracker

    def test_track_response_basic_tokens(self) -> None:
        wrapper, tracker = self._make_wrapper()
        mock_response = MagicMock()
        mock_response.model = "claude-3-5-sonnet"
        mock_response.usage.input_tokens = 200
        mock_response.usage.output_tokens = 100
        mock_response.usage.cache_read_input_tokens = None
        mock_response.usage.cache_creation_input_tokens = None
        del mock_response.usage.thinking_input_tokens
        del mock_response.usage.thinking_tokens

        with patch.object(tracker, "track") as mock_track:
            wrapper._track_response(mock_response)
            mock_track.assert_called_once_with(
                provider="anthropic",
                model="claude-3-5-sonnet",
                input_tokens=200,
                output_tokens=100,
                cached_read_tokens=None,
                cached_write_tokens=None,
                thinking_tokens=None,
            )

        tracker.close()

    def test_track_response_thinking_zero_preserved(self) -> None:
        wrapper, tracker = self._make_wrapper()
        mock_response = MagicMock()
        mock_response.model = "claude-3-7-sonnet"
        mock_response.usage.input_tokens = 100
        mock_response.usage.output_tokens = 50
        mock_response.usage.cache_read_input_tokens = None
        mock_response.usage.cache_creation_input_tokens = None
        mock_response.usage.thinking_input_tokens = 0
        mock_response.usage.thinking_tokens = 400

        with patch.object(tracker, "track") as mock_track:
            wrapper._track_response(mock_response)
            call_kwargs = mock_track.call_args[1]
            assert call_kwargs["thinking_tokens"] == 0

        tracker.close()

    def test_wrap_stream_tracks_on_partial_consumption(self) -> None:
        wrapper, tracker = self._make_wrapper()

        message_start = MagicMock()
        message_start.type = "message_start"
        message_start.message.usage.input_tokens = 150
        message_start.message.usage.cache_read_input_tokens = None
        message_start.message.usage.cache_creation_input_tokens = None
        message_start.message.usage.thinking_input_tokens = None
        message_start.message.usage.thinking_tokens = None

        message_delta = MagicMock()
        message_delta.type = "message_delta"
        message_delta.usage.output_tokens = 75

        with patch.object(tracker, "track") as mock_track:
            wrapped = wrapper._wrap_stream(
                iter([message_start, message_delta]), "claude-3-5-sonnet"
            )
            next(wrapped)
            wrapped_generator: Any = wrapped
            wrapped_generator.close()
            mock_track.assert_called_once()
            call_kwargs = mock_track.call_args[1]
            assert call_kwargs["input_tokens"] == 150

        tracker.close()


class TestOpenAICompletionsWrapper:
    def _make_wrapper(
        self, tracker: TokpulseTracker | None = None
    ) -> tuple[OpenAICompletionsWrapper, TokpulseTracker]:
        if tracker is None:
            tracker = _mock_tracker()
        mock_completions = MagicMock()
        return OpenAICompletionsWrapper(mock_completions, tracker), tracker

    def test_wrap_stream_tracks_on_partial_consumption(self) -> None:
        wrapper, tracker = self._make_wrapper()

        chunk_with_usage = MagicMock()
        chunk_with_usage.usage.prompt_tokens = 42
        chunk_with_usage.usage.completion_tokens = 17
        chunk_with_usage.usage.prompt_tokens_details = None

        with patch.object(tracker, "track") as mock_track:
            wrapped = wrapper._wrap_stream(iter([chunk_with_usage, MagicMock()]), "gpt-4o")
            next(wrapped)
            wrapped_generator: Any = wrapped
            wrapped_generator.close()
            mock_track.assert_called_once()
            call_kwargs = mock_track.call_args[1]
            assert call_kwargs["input_tokens"] == 42
            assert call_kwargs["output_tokens"] == 17

        tracker.close()


class TestGoogleWrapper:
    def _make_wrapper(
        self, tracker: TokpulseTracker | None = None
    ) -> tuple[GoogleWrapper, TokpulseTracker]:
        if tracker is None:
            tracker = _mock_tracker()

        fake_model = MagicMock()
        fake_model.model_name = "gemini-2.5-flash"
        return GoogleWrapper(fake_model, tracker), tracker

    def test_wrap_stream_tracks_on_partial_consumption(self) -> None:
        wrapper, tracker = self._make_wrapper()

        chunk = MagicMock()
        chunk.usage_metadata.prompt_token_count = 80
        chunk.usage_metadata.candidates_token_count = 25
        chunk.usage_metadata.cached_content_token_count = 10

        with patch.object(tracker, "track") as mock_track:
            wrapped = wrapper._wrap_stream(iter([chunk, MagicMock()]))
            next(wrapped)
            wrapped_generator: Any = wrapped
            wrapped_generator.close()
            mock_track.assert_called_once()
            call_kwargs = mock_track.call_args[1]
            assert call_kwargs["input_tokens"] == 80
            assert call_kwargs["output_tokens"] == 25

        tracker.close()


# ---------------------------------------------------------------------------
# Async wrapper tests
# ---------------------------------------------------------------------------


async def _aiter(items: list[Any]) -> Any:
    """Yield items from a list as an async iterator."""
    for item in items:
        yield item


class TestAsyncOpenAICompletionsWrapper:
    def _make_wrapper(
        self, tracker: TokpulseTracker | None = None
    ) -> tuple[AsyncOpenAICompletionsWrapper, TokpulseTracker]:
        if tracker is None:
            tracker = _mock_tracker()
        mock_completions = MagicMock()
        return AsyncOpenAICompletionsWrapper(mock_completions, tracker), tracker

    async def test_non_streaming_tracks_tokens(self) -> None:
        wrapper, tracker = self._make_wrapper()
        mock_response = MagicMock()
        mock_response.model = "gpt-4o"
        mock_response.usage.prompt_tokens = 100
        mock_response.usage.completion_tokens = 50
        mock_response.usage.prompt_tokens_details = None
        wrapper._completions.create = AsyncMock(return_value=mock_response)

        with patch.object(tracker, "track") as mock_track:
            await wrapper.create(model="gpt-4o", messages=[])
            mock_track.assert_called_once_with(
                provider="openai",
                model="gpt-4o",
                input_tokens=100,
                output_tokens=50,
                cached_input_tokens=None,
            )

        tracker.close()

    async def test_non_streaming_no_usage_skips_track(self) -> None:
        wrapper, tracker = self._make_wrapper()
        mock_response = MagicMock(spec=[])  # no usage attribute
        wrapper._completions.create = AsyncMock(return_value=mock_response)

        with patch.object(tracker, "track") as mock_track:
            await wrapper.create(model="gpt-4o", messages=[])
            mock_track.assert_not_called()

        tracker.close()

    # BOUNDARY: usage present but None
    async def test_non_streaming_usage_none_skips_track(self) -> None:
        wrapper, tracker = self._make_wrapper()
        mock_response = MagicMock()
        mock_response.usage = None
        wrapper._completions.create = AsyncMock(return_value=mock_response)

        with patch.object(tracker, "track") as mock_track:
            await wrapper.create(model="gpt-4o", messages=[])
            mock_track.assert_not_called()

        tracker.close()

    async def test_streaming_tracks_tokens_from_usage_chunk(self) -> None:
        wrapper, tracker = self._make_wrapper()
        chunk_with_usage = MagicMock()
        chunk_with_usage.usage.prompt_tokens = 30
        chunk_with_usage.usage.completion_tokens = 15
        chunk_with_usage.usage.prompt_tokens_details = None
        wrapper._completions.create = AsyncMock(return_value=_aiter([chunk_with_usage]))

        with patch.object(tracker, "track") as mock_track:
            gen = await wrapper.create(model="gpt-4o", stream=True, messages=[])
            async for _ in gen:
                pass
            mock_track.assert_called_once()
            call_kwargs = mock_track.call_args[1]
            assert call_kwargs["input_tokens"] == 30
            assert call_kwargs["output_tokens"] == 15

        tracker.close()

    async def test_streaming_tracks_on_partial_consumption(self) -> None:
        wrapper, tracker = self._make_wrapper()
        chunk = MagicMock()
        chunk.usage.prompt_tokens = 20
        chunk.usage.completion_tokens = 10
        chunk.usage.prompt_tokens_details = None
        wrapper._completions.create = AsyncMock(return_value=_aiter([chunk, MagicMock()]))

        with patch.object(tracker, "track") as mock_track:
            gen = await wrapper.create(model="gpt-4o", stream=True, messages=[])
            async for _ in gen:
                break  # partial consumption
            # Async generator finally runs only after aclose() is awaited.
            await gen.aclose()
            mock_track.assert_called_once()

        tracker.close()


class TestAsyncAnthropicMessagesWrapper:
    def _make_wrapper(
        self, tracker: TokpulseTracker | None = None
    ) -> tuple[AsyncAnthropicMessagesWrapper, TokpulseTracker]:
        if tracker is None:
            tracker = _mock_tracker()
        mock_messages = MagicMock()
        return AsyncAnthropicMessagesWrapper(mock_messages, tracker), tracker

    async def test_non_streaming_tracks_tokens(self) -> None:
        wrapper, tracker = self._make_wrapper()
        mock_response = MagicMock()
        mock_response.model = "claude-3-5-sonnet"
        mock_response.usage.input_tokens = 200
        mock_response.usage.output_tokens = 100
        mock_response.usage.cache_read_input_tokens = None
        mock_response.usage.cache_creation_input_tokens = None
        del mock_response.usage.thinking_input_tokens
        del mock_response.usage.thinking_tokens
        wrapper._messages.create = AsyncMock(return_value=mock_response)

        with patch.object(tracker, "track") as mock_track:
            await wrapper.create(model="claude-3-5-sonnet", max_tokens=512, messages=[])
            mock_track.assert_called_once_with(
                provider="anthropic",
                model="claude-3-5-sonnet",
                input_tokens=200,
                output_tokens=100,
                cached_read_tokens=None,
                cached_write_tokens=None,
                thinking_tokens=None,
            )

        tracker.close()

    # BOUNDARY: usage present but None
    async def test_non_streaming_usage_none_skips_track(self) -> None:
        wrapper, tracker = self._make_wrapper()
        mock_response = MagicMock()
        mock_response.usage = None
        wrapper._messages.create = AsyncMock(return_value=mock_response)

        with patch.object(tracker, "track") as mock_track:
            await wrapper.create(model="claude-3-5-sonnet", max_tokens=512, messages=[])
            mock_track.assert_not_called()

        tracker.close()

    async def test_streaming_tracks_tokens(self) -> None:
        wrapper, tracker = self._make_wrapper()
        message_start = MagicMock()
        message_start.type = "message_start"
        message_start.message.usage.input_tokens = 150
        message_start.message.usage.cache_read_input_tokens = None
        message_start.message.usage.cache_creation_input_tokens = None
        message_start.message.usage.thinking_input_tokens = None
        message_start.message.usage.thinking_tokens = None
        message_delta = MagicMock()
        message_delta.type = "message_delta"
        message_delta.usage.output_tokens = 60
        wrapper._messages.create = AsyncMock(return_value=_aiter([message_start, message_delta]))

        with patch.object(tracker, "track") as mock_track:
            gen = await wrapper.create(
                model="claude-3-5-sonnet", stream=True, max_tokens=512, messages=[]
            )
            async for _ in gen:
                pass
            mock_track.assert_called_once()
            call_kwargs = mock_track.call_args[1]
            assert call_kwargs["input_tokens"] == 150
            assert call_kwargs["output_tokens"] == 60

        tracker.close()

    async def test_streaming_tracks_on_partial_consumption(self) -> None:
        wrapper, tracker = self._make_wrapper()
        message_start = MagicMock()
        message_start.type = "message_start"
        message_start.message.usage.input_tokens = 80
        message_start.message.usage.cache_read_input_tokens = None
        message_start.message.usage.cache_creation_input_tokens = None
        message_start.message.usage.thinking_input_tokens = None
        message_start.message.usage.thinking_tokens = None
        wrapper._messages.create = AsyncMock(return_value=_aiter([message_start]))

        with patch.object(tracker, "track") as mock_track:
            gen = await wrapper.create(
                model="claude-3-5-sonnet", stream=True, max_tokens=512, messages=[]
            )
            async for _ in gen:
                break  # partial consumption
            # Async generator finally runs only after aclose() is awaited.
            await gen.aclose()
            mock_track.assert_called_once()
            call_kwargs = mock_track.call_args[1]
            assert call_kwargs["input_tokens"] == 80

        tracker.close()

    async def test_thinking_tokens_tracked(self) -> None:
        wrapper, tracker = self._make_wrapper()
        mock_response = MagicMock()
        mock_response.model = "claude-3-7-sonnet"
        mock_response.usage.input_tokens = 100
        mock_response.usage.output_tokens = 50
        mock_response.usage.cache_read_input_tokens = None
        mock_response.usage.cache_creation_input_tokens = None
        mock_response.usage.thinking_input_tokens = 300
        del mock_response.usage.thinking_tokens
        wrapper._messages.create = AsyncMock(return_value=mock_response)

        with patch.object(tracker, "track") as mock_track:
            await wrapper.create(model="claude-3-7-sonnet", max_tokens=1024, messages=[])
            call_kwargs = mock_track.call_args[1]
            assert call_kwargs["thinking_tokens"] == 300

        tracker.close()


class TestAsyncWrapperFactory:
    def test_openai_async_returns_async_wrapper(self) -> None:
        tracker = _mock_tracker()
        fake_client = MagicMock()
        wrapped = tracker.openai_async(fake_client)
        assert isinstance(wrapped, AsyncOpenAIWrapper)
        tracker.close()

    def test_anthropic_async_returns_async_wrapper(self) -> None:
        tracker = _mock_tracker()
        fake_client = MagicMock()
        wrapped = tracker.anthropic_async(fake_client)
        assert isinstance(wrapped, AsyncAnthropicWrapper)
        tracker.close()


# ---------------------------------------------------------------------------
# messages.stream() context manager — sync
# ---------------------------------------------------------------------------


class TestAnthropicMessagesStream:
    def _make_wrapper(
        self, tracker: TokpulseTracker | None = None
    ) -> tuple[AnthropicMessagesWrapper, TokpulseTracker]:
        if tracker is None:
            tracker = _mock_tracker()
        mock_messages = MagicMock()
        return AnthropicMessagesWrapper(mock_messages, tracker), tracker

    def _make_final_message(
        self,
        *,
        model: str = "claude-3-5-sonnet",
        input_tokens: int = 200,
        output_tokens: int = 80,
        cache_read: int | None = None,
        cache_write: int | None = None,
        thinking: int | None = None,
    ) -> MagicMock:
        msg = MagicMock()
        msg.model = model
        msg.usage.input_tokens = input_tokens
        msg.usage.output_tokens = output_tokens
        msg.usage.cache_read_input_tokens = cache_read
        msg.usage.cache_creation_input_tokens = cache_write
        if thinking is not None:
            msg.usage.thinking_input_tokens = thinking
        else:
            del msg.usage.thinking_input_tokens
            del msg.usage.thinking_tokens
        return msg

    def test_stream_tracks_on_context_exit(self) -> None:
        wrapper, tracker = self._make_wrapper()
        final_msg = self._make_final_message(input_tokens=300, output_tokens=120)

        mock_stream = MagicMock()
        mock_stream.get_final_message.return_value = final_msg
        mock_stream_mgr = MagicMock()
        mock_stream_mgr.__enter__ = MagicMock(return_value=mock_stream)
        mock_stream_mgr.__exit__ = MagicMock(return_value=None)
        wrapper._messages.stream.return_value = mock_stream_mgr

        with patch.object(tracker, "track") as mock_track:
            with wrapper.stream(model="claude-3-5-sonnet", max_tokens=512, messages=[]) as s:
                assert s is mock_stream

            mock_track.assert_called_once_with(
                provider="anthropic",
                model="claude-3-5-sonnet",
                input_tokens=300,
                output_tokens=120,
                cached_read_tokens=None,
                cached_write_tokens=None,
                thinking_tokens=None,
            )
        tracker.close()

    def test_stream_tracks_cache_and_thinking_tokens(self) -> None:
        wrapper, tracker = self._make_wrapper()
        final_msg = self._make_final_message(
            input_tokens=400,
            output_tokens=150,
            cache_read=50,
            cache_write=20,
            thinking=500,
        )

        mock_stream = MagicMock()
        mock_stream.get_final_message.return_value = final_msg
        mock_stream_mgr = MagicMock()
        mock_stream_mgr.__enter__ = MagicMock(return_value=mock_stream)
        mock_stream_mgr.__exit__ = MagicMock(return_value=None)
        wrapper._messages.stream.return_value = mock_stream_mgr

        with patch.object(tracker, "track") as mock_track:
            with wrapper.stream(model="claude-3-7-sonnet", max_tokens=1024, messages=[]):
                pass
            call_kwargs = mock_track.call_args[1]
            assert call_kwargs["cached_read_tokens"] == 50
            assert call_kwargs["cached_write_tokens"] == 20
            assert call_kwargs["thinking_tokens"] == 500
        tracker.close()

    def test_stream_silently_skips_track_when_get_final_message_raises(self) -> None:
        """If the stream was not fully consumed, get_final_message() may raise; no crash."""
        wrapper, tracker = self._make_wrapper()

        mock_stream = MagicMock()
        mock_stream.get_final_message.side_effect = RuntimeError("Stream not complete")
        mock_stream_mgr = MagicMock()
        mock_stream_mgr.__enter__ = MagicMock(return_value=mock_stream)
        mock_stream_mgr.__exit__ = MagicMock(return_value=None)
        wrapper._messages.stream.return_value = mock_stream_mgr

        with patch.object(tracker, "track") as mock_track:
            with wrapper.stream(model="claude-3-5-sonnet", max_tokens=512, messages=[]):
                pass
            mock_track.assert_not_called()
        tracker.close()

    def test_stream_returns_tracked_stream_manager(self) -> None:
        wrapper, tracker = self._make_wrapper()
        from tokpulse.wrappers import _TrackedStreamManager

        mock_stream_mgr = MagicMock()
        wrapper._messages.stream.return_value = mock_stream_mgr

        result = wrapper.stream(model="claude-3-5-sonnet", max_tokens=512, messages=[])
        assert isinstance(result, _TrackedStreamManager)
        tracker.close()


# ---------------------------------------------------------------------------
# client.beta.messages — sync
# ---------------------------------------------------------------------------


class TestAnthropicBetaWrapper:
    def test_beta_property_returns_beta_wrapper(self) -> None:
        tracker = _mock_tracker()
        fake_client = MagicMock()
        wrapped = AnthropicWrapper(fake_client, tracker)
        assert isinstance(wrapped.beta, AnthropicBetaWrapper)
        tracker.close()

    def test_beta_messages_returns_messages_wrapper(self) -> None:
        tracker = _mock_tracker()
        fake_client = MagicMock()
        beta_wrapper = AnthropicBetaWrapper(fake_client.beta, tracker)
        assert isinstance(beta_wrapper.messages, AnthropicMessagesWrapper)
        tracker.close()

    def test_beta_messages_create_tracks_usage(self) -> None:
        tracker = _mock_tracker()
        fake_client = MagicMock()
        wrapped = AnthropicWrapper(fake_client, tracker)

        mock_response = MagicMock()
        mock_response.model = "claude-3-5-sonnet"
        mock_response.usage.input_tokens = 100
        mock_response.usage.output_tokens = 50
        mock_response.usage.cache_read_input_tokens = None
        mock_response.usage.cache_creation_input_tokens = None
        del mock_response.usage.thinking_input_tokens
        del mock_response.usage.thinking_tokens
        fake_client.beta.messages.create.return_value = mock_response

        with patch.object(tracker, "track") as mock_track:
            wrapped.beta.messages.create(model="claude-3-5-sonnet", max_tokens=512, messages=[])
            mock_track.assert_called_once_with(
                provider="anthropic",
                model="claude-3-5-sonnet",
                input_tokens=100,
                output_tokens=50,
                cached_read_tokens=None,
                cached_write_tokens=None,
                thinking_tokens=None,
            )
        tracker.close()

    def test_beta_proxies_non_messages_attributes(self) -> None:
        tracker = _mock_tracker()
        fake_client = MagicMock()
        fake_client.beta.tools = "tools_sentinel"
        beta_wrapper = AnthropicBetaWrapper(fake_client.beta, tracker)
        assert beta_wrapper.tools == "tools_sentinel"
        tracker.close()


# ---------------------------------------------------------------------------
# messages.stream() context manager — async
# ---------------------------------------------------------------------------


class TestAsyncAnthropicMessagesStream:
    def _make_wrapper(
        self, tracker: TokpulseTracker | None = None
    ) -> tuple[AsyncAnthropicMessagesWrapper, TokpulseTracker]:
        if tracker is None:
            tracker = _mock_tracker()
        mock_messages = MagicMock()
        return AsyncAnthropicMessagesWrapper(mock_messages, tracker), tracker

    async def test_stream_tracks_on_context_exit(self) -> None:
        wrapper, tracker = self._make_wrapper()

        final_msg = MagicMock()
        final_msg.model = "claude-3-5-sonnet"
        final_msg.usage.input_tokens = 200
        final_msg.usage.output_tokens = 90
        final_msg.usage.cache_read_input_tokens = None
        final_msg.usage.cache_creation_input_tokens = None
        del final_msg.usage.thinking_input_tokens
        del final_msg.usage.thinking_tokens

        mock_stream = MagicMock()
        mock_stream.get_final_message.return_value = final_msg
        mock_stream_mgr = AsyncMock()
        mock_stream_mgr.__aenter__ = AsyncMock(return_value=mock_stream)
        mock_stream_mgr.__aexit__ = AsyncMock(return_value=None)
        wrapper._messages.stream.return_value = mock_stream_mgr

        with patch.object(tracker, "track") as mock_track:
            async with wrapper.stream(model="claude-3-5-sonnet", max_tokens=512, messages=[]) as s:
                assert s is mock_stream

            mock_track.assert_called_once_with(
                provider="anthropic",
                model="claude-3-5-sonnet",
                input_tokens=200,
                output_tokens=90,
                cached_read_tokens=None,
                cached_write_tokens=None,
                thinking_tokens=None,
            )
        tracker.close()

    async def test_stream_silently_skips_when_get_final_message_raises(self) -> None:
        wrapper, tracker = self._make_wrapper()

        mock_stream = MagicMock()
        mock_stream.get_final_message.side_effect = RuntimeError("not done")
        mock_stream_mgr = AsyncMock()
        mock_stream_mgr.__aenter__ = AsyncMock(return_value=mock_stream)
        mock_stream_mgr.__aexit__ = AsyncMock(return_value=None)
        wrapper._messages.stream.return_value = mock_stream_mgr

        with patch.object(tracker, "track") as mock_track:
            async with wrapper.stream(model="claude-3-5-sonnet", max_tokens=512, messages=[]):
                pass
            mock_track.assert_not_called()
        tracker.close()

    async def test_stream_returns_async_tracked_stream_manager(self) -> None:
        wrapper, tracker = self._make_wrapper()
        from tokpulse.wrappers import _AsyncTrackedStreamManager

        mock_stream_mgr = MagicMock()
        wrapper._messages.stream.return_value = mock_stream_mgr

        result = wrapper.stream(model="claude-3-5-sonnet", max_tokens=512, messages=[])
        assert isinstance(result, _AsyncTrackedStreamManager)
        tracker.close()


# ---------------------------------------------------------------------------
# client.beta.messages — async
# ---------------------------------------------------------------------------


class TestAsyncAnthropicBetaWrapper:
    def test_beta_property_returns_async_beta_wrapper(self) -> None:
        tracker = _mock_tracker()
        fake_client = MagicMock()
        wrapped = AsyncAnthropicWrapper(fake_client, tracker)
        assert isinstance(wrapped.beta, AsyncAnthropicBetaWrapper)
        tracker.close()

    def test_beta_messages_returns_async_messages_wrapper(self) -> None:
        tracker = _mock_tracker()
        fake_client = MagicMock()
        beta_wrapper = AsyncAnthropicBetaWrapper(fake_client.beta, tracker)
        assert isinstance(beta_wrapper.messages, AsyncAnthropicMessagesWrapper)
        tracker.close()

    async def test_beta_messages_create_tracks_usage(self) -> None:
        tracker = _mock_tracker()
        fake_client = MagicMock()
        wrapped = AsyncAnthropicWrapper(fake_client, tracker)

        mock_response = MagicMock()
        mock_response.model = "claude-3-5-sonnet"
        mock_response.usage.input_tokens = 150
        mock_response.usage.output_tokens = 70
        mock_response.usage.cache_read_input_tokens = None
        mock_response.usage.cache_creation_input_tokens = None
        del mock_response.usage.thinking_input_tokens
        del mock_response.usage.thinking_tokens
        fake_client.beta.messages.create = AsyncMock(return_value=mock_response)

        with patch.object(tracker, "track") as mock_track:
            await wrapped.beta.messages.create(
                model="claude-3-5-sonnet", max_tokens=512, messages=[]
            )
            mock_track.assert_called_once_with(
                provider="anthropic",
                model="claude-3-5-sonnet",
                input_tokens=150,
                output_tokens=70,
                cached_read_tokens=None,
                cached_write_tokens=None,
                thinking_tokens=None,
            )
        tracker.close()
