⚡️ Speed up method APIRequestor._interpret_async_response by 11% - #13

import asyncio  # used to run async functions
# --- Function under test (as provided, do not modify) ---
import json
from unittest.mock import AsyncMock, MagicMock, patch

import pytest  # used for our unit tests
from together.abstract.api_requestor import APIRequestor


class DummyTogetherResponse:
    def __init__(self, data, headers):
        self.data = data
        self._headers = headers
        self.request_id = headers.get("X-Request-ID", None)

class DummyError(Exception):
    pass

class DummyTimeout(DummyError):
    pass

class DummyAPIError(DummyError):
    pass

class DummyRateLimitError(DummyError):
    pass

class DummyInvalidRequestError(DummyError):
    pass

class DummyAuthenticationError(DummyError):
    pass

class DummyServiceUnavailableError(DummyError):
    pass

# Dummy TogetherClient
class DummyTogetherClient:
    def __init__(self, base_url=None, api_key=None, max_retries=None, supplied_headers=None, timeout=None):
        self.base_url = base_url
        self.api_key = api_key
        self.max_retries = max_retries
        self.supplied_headers = supplied_headers or {}
        self.timeout = timeout

# Dummy TogetherResponse
TogetherResponse = DummyTogetherResponse

# --- END function under test ---

# --- Test suite for APIRequestor._interpret_async_response ---

@pytest.fixture
def api_requestor():
    # Provide a dummy TogetherClient for the APIRequestor
    client = DummyTogetherClient(
        base_url="https://api.test/",
        api_key="testkey",
        max_retries=2,
        supplied_headers={"X-Test": "yes"},
        timeout=5
    )
    return APIRequestor(client)

# Helper: Dummy aiohttp.ClientResponse-like object for non-streaming
class DummyAiohttpResponse:
    def __init__(self, body, status=200, headers=None, raise_on_read=None):
        self._body = body
        self.status = status
        self.headers = headers or {"Content-Type": "application/json"}
        self._read_called = False
        self._raise_on_read = raise_on_read
        self.content = None  # For streaming, set by tests

    async def read(self):
        if self._raise_on_read:
            raise self._raise_on_read
        self._read_called = True
        if isinstance(self._body, str):
            return self._body.encode("utf-8")
        return self._body

# Helper: Dummy async iterator for streaming content
class DummyAsyncStream:
    def __init__(self, lines):
        self._lines = lines
        self._idx = 0

    def __aiter__(self):
        return self

    async def __anext__(self):
        if self._idx >= len(self._lines):
            raise StopAsyncIteration
        val = self._lines[self._idx]
        self._idx += 1
        await asyncio.sleep(0)  # yield control to event loop
        return val

# ---- Basic Test Cases ----

@pytest.mark.asyncio
async def test_interpret_async_response_basic_json(api_requestor):
    """Test normal JSON response, non-stream."""
    body = '{"foo": "bar"}'
    headers = {"Content-Type": "application/json"}
    resp = DummyAiohttpResponse(body, status=200, headers=headers)
    result, is_stream = await api_requestor._interpret_async_response(resp, stream=False)

@pytest.mark.asyncio
async def test_interpret_async_response_basic_plaintext(api_requestor):
    """Test plain text response, non-stream."""
    body = "hello world"
    headers = {"Content-Type": "text/plain"}
    resp = DummyAiohttpResponse(body, status=200, headers=headers)
    result, is_stream = await api_requestor._interpret_async_response(resp, stream=False)

@pytest.mark.asyncio
async def test_interpret_async_response_basic_204(api_requestor):
    """Test HTTP 204 (no content) response."""
    body = ""
    headers = {"Content-Type": "application/json"}
    resp = DummyAiohttpResponse(body, status=204, headers=headers)
    result, is_stream = await api_requestor._interpret_async_response(resp, stream=False)

# ---- Edge Test Cases ----

@pytest.mark.asyncio





async def test_interpret_async_response_timeout(api_requestor):
    """Test that TimeoutError is raised on read timeout."""
    headers = {"Content-Type": "application/json"}
    resp = DummyAiohttpResponse("{}", status=200, headers=headers, raise_on_read=DummyTimeout())
    with pytest.raises(DummyTimeout):
        await api_requestor._interpret_async_response(resp, stream=False)

@pytest.mark.asyncio




async def test_interpret_async_response_concurrent_non_stream(api_requestor):
    """Test concurrent non-streaming responses."""
    bodies = [json.dumps({"i": i}) for i in range(10)]
    headers = {"Content-Type": "application/json"}
    resps = [DummyAiohttpResponse(body, status=200, headers=headers) for body in bodies]
    coros = [api_requestor._interpret_async_response(resp, stream=False) for resp in resps]
    results = await asyncio.gather(*coros)
    for idx, (result, is_stream) in enumerate(results):
        pass

@pytest.mark.asyncio

async def test_interpret_async_response_throughput_small_load(api_requestor):
    """Throughput: small batch of non-streaming requests."""
    bodies = [json.dumps({"val": i}) for i in range(5)]
    headers = {"Content-Type": "application/json"}
    resps = [DummyAiohttpResponse(body, status=200, headers=headers) for body in bodies]
    coros = [api_requestor._interpret_async_response(resp, stream=False) for resp in resps]
    results = await asyncio.gather(*coros)

@pytest.mark.asyncio

async def test_interpret_async_response_throughput_high_volume(api_requestor):
    """Throughput: high volume of concurrent non-streaming requests."""
    # Keep under 1000 to avoid excessive test time
    N = 100
    bodies = [json.dumps({"num": i}) for i in range(N)]
    headers = {"Content-Type": "application/json"}
    resps = [DummyAiohttpResponse(body, status=200, headers=headers) for body in bodies]
    coros = [api_requestor._interpret_async_response(resp, stream=False) for resp in resps]
    results = await asyncio.gather(*coros)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
import asyncio  # used to run async functions
import json
import sys
from types import SimpleNamespace
from typing import AsyncGenerator

import pytest  # used for our unit tests
from together import error, utils
from together.abstract.api_requestor import APIRequestor
from together.constants import BASE_URL, MAX_RETRIES, TIMEOUT_SECS
from together.together_response import TogetherResponse
from together.types import TogetherClient
from together.types.error import TogetherErrorResponse

# ---- Mocks and helpers for aiohttp and TogetherResponse ----

class MockStreamReader:
    """A mock of aiohttp.StreamReader for async iteration."""
    def __init__(self, lines):
        self._lines = lines

    def __aiter__(self):
        self._iter = iter(self._lines)
        return self

    async def __anext__(self):
        try:
            return next(self._iter)
        except StopIteration:
            raise StopAsyncIteration

class MockClientResponse:
    """A mock of aiohttp.ClientResponse."""
    def __init__(self, *, content=None, headers=None, status=200, body=b"", raise_on_read=None):
        self.content = content
        self.headers = headers or {}
        self.status = status
        self._body = body
        self._read_called = False
        self._raise_on_read = raise_on_read

    async def read(self):
        if self._raise_on_read:
            raise self._raise_on_read
        self._read_called = True
        return self._body

class MockTogetherResponse:
    """A mock of TogetherResponse."""
    def __init__(self, data, headers):
        self.data = data
        self._headers = headers
        self.request_id = headers.get("X-Request-Id", "req_123")

# ---- Patch the Together imports ----


# Patch together.error
class DummyError(Exception): pass
class DummyTimeout(DummyError): pass
class DummyAPIError(DummyError): pass
class DummyAuthenticationError(DummyError): pass
class DummyInvalidRequestError(DummyError): pass
class DummyRateLimitError(DummyError): pass
class DummyServiceUnavailableError(DummyError): pass
TogetherResponse = MockTogetherResponse

# Patch together.types
class DummyClient:
    def __init__(self, base_url=None, api_key=None, max_retries=None, supplied_headers=None, timeout=None):
        self.base_url = base_url
        self.api_key = api_key
        self.max_retries = max_retries
        self.supplied_headers = supplied_headers or {}
        self.timeout = timeout

# ---- TESTS ----

@pytest.mark.asyncio