test_request.py

from __future__ import annotations

import asyncio
import io
import json
from collections.abc import Iterator
from http import HTTPStatus
from typing import TYPE_CHECKING, Any
from unittest import mock
from unittest.mock import AsyncMock

import aiohttp
import pytest
from aioresponses import aioresponses

from ai.backend.client.config import API_VERSION, get_config
from ai.backend.client.exceptions import BackendAPIError, BackendClientError
from ai.backend.client.request import AttachedFile, Request, Response
from ai.backend.client.session import AsyncSession, Session

if TYPE_CHECKING:
    from ai.backend.client.config import APIConfig


@pytest.fixture(scope="module", autouse=True)
def api_version() -> Iterator[None]:
    mock_nego_func = AsyncMock()
    mock_nego_func.return_value = API_VERSION
    with mock.patch("ai.backend.client.session._negotiate_api_version", mock_nego_func):
        yield


@pytest.fixture
def session(defconfig: APIConfig) -> Iterator[Session]:
    with Session(config=defconfig) as session:
        yield session


@pytest.fixture
def mock_request_params(session: Session) -> dict[str, Any]:
    return {
        "method": "GET",
        "path": "/function/item/",
        "params": {"app": "999"},
        "content": b'{"test1": 1}',
        "content_type": "application/json",
    }


def test_request_initialization(mock_request_params: dict[str, Any]) -> None:
    rqst = Request(**mock_request_params)

    assert rqst.method == mock_request_params["method"]
    assert rqst.params == mock_request_params["params"]
    assert rqst.path == mock_request_params["path"].lstrip("/")
    assert rqst.content == mock_request_params["content"]
    assert "X-BackendAI-Version" in rqst.headers


def test_request_set_content_none(mock_request_params: dict[str, Any]) -> None:
    mock_request_params = mock_request_params.copy()
    mock_request_params["content"] = None
    rqst = Request(**mock_request_params)
    assert rqst.content == b""
    assert rqst._pack_content() is rqst.content


def test_request_set_content(mock_request_params: dict[str, Any]) -> None:
    rqst = Request(**mock_request_params)
    assert rqst.content == mock_request_params["content"]
    assert rqst.content_type == "application/json"
    assert rqst._pack_content() is rqst.content

    mock_request_params["content"] = "hello"
    mock_request_params["content_type"] = None
    rqst = Request(**mock_request_params)
    assert rqst.content == b"hello"
    assert rqst.content_type == "text/plain"
    assert rqst._pack_content() is rqst.content

    mock_request_params["content"] = b"\x00\x01\xfe\xff"
    mock_request_params["content_type"] = None
    rqst = Request(**mock_request_params)
    assert rqst.content == b"\x00\x01\xfe\xff"
    assert rqst.content_type == "application/octet-stream"
    assert rqst._pack_content() is rqst.content


def test_request_attach_files(mock_request_params: dict[str, Any]) -> None:
    files = [
        AttachedFile("test1.txt", io.BytesIO(), "application/octet-stream"),
        AttachedFile("test2.txt", io.BytesIO(), "application/octet-stream"),
    ]

    mock_request_params["content"] = b"something"
    rqst = Request(**mock_request_params)
    with pytest.raises(ValueError):
        rqst.attach_files(files)

    mock_request_params["content"] = b""
    rqst = Request(**mock_request_params)
    rqst.attach_files(files)

    assert rqst.content_type == "multipart/form-data"
    assert rqst.content == b""
    packed_content = rqst._pack_content()
    assert isinstance(packed_content, aiohttp.FormData)
    assert packed_content.is_multipart


def test_build_correct_url(mock_request_params: dict[str, Any]) -> None:
    config = get_config()
    canonical_url = str(config.endpoint).rstrip("/") + "/function?app=999"

    mock_request_params["path"] = "/function"
    rqst = Request(**mock_request_params)
    assert str(rqst._build_url()) == canonical_url

    mock_request_params["path"] = "function"
    rqst = Request(**mock_request_params)
    assert str(rqst._build_url()) == canonical_url


async def test_fetch_invalid_method(mock_request_params: dict[str, Any]) -> None:
    mock_request_params["method"] = "STRANGE"
    rqst = Request(**mock_request_params)

    with pytest.raises(ValueError):
        async with rqst.fetch():
            pass


async def test_fetch(dummy_endpoint: str) -> None:
    with aioresponses() as m, Session():
        body = b"hello world"
        m.post(
            dummy_endpoint + "function",
            status=HTTPStatus.OK,
            body=body,
            headers={"Content-Type": "text/plain; charset=utf-8", "Content-Length": str(len(body))},
        )
        rqst = Request("POST", "function")
        async with rqst.fetch() as resp:
            assert isinstance(resp, Response)
            assert resp.status == HTTPStatus.OK
            assert resp.content_type == "text/plain"
            assert await resp.text() == body.decode()
            assert resp.content_length == len(body)

    with aioresponses() as m, Session():
        body = b'{"a": 1234, "b": null}'
        m.post(
            dummy_endpoint + "function",
            status=HTTPStatus.OK,
            body=body,
            headers={
                "Content-Type": "application/json; charset=utf-8",
                "Content-Length": str(len(body)),
            },
        )
        rqst = Request("POST", "function")
        async with rqst.fetch() as resp:
            assert isinstance(resp, Response)
            assert resp.status == HTTPStatus.OK
            assert resp.content_type == "application/json"
            assert await resp.text() == body.decode()
            assert await resp.json() == {"a": 1234, "b": None}
            assert resp.content_length == len(body)


async def test_streaming_fetch(dummy_endpoint: str) -> None:
    # Read content by chunks.
    with aioresponses() as m, Session():
        body = b"hello world"
        m.post(
            dummy_endpoint + "function",
            status=HTTPStatus.OK,
            body=body,
            headers={"Content-Type": "text/plain; charset=utf-8", "Content-Length": str(len(body))},
        )
        rqst = Request("POST", "function")
        async with rqst.fetch() as resp:
            assert resp.status == HTTPStatus.OK
            assert resp.content_type == "text/plain"
            assert await resp.read(3) == b"hel"
            assert await resp.read(2) == b"lo"
            await resp.read()
            with pytest.raises(AssertionError):
                assert await resp.text()


async def test_invalid_requests(dummy_endpoint: str) -> None:
    with aioresponses() as m, Session():
        body = json.dumps({
            "type": "https://api.backend.ai/probs/kernel-not-found",
            "title": "Kernel Not Found",
        }).encode("utf8")
        m.post(
            dummy_endpoint,
            status=HTTPStatus.NOT_FOUND,
            body=body,
            headers={
                "Content-Type": "application/problem+json; charset=utf-8",
                "Content-Length": str(len(body)),
            },
        )
        rqst = Request("POST", "/")
        with pytest.raises(BackendAPIError) as e:
            async with rqst.fetch():
                pass
            assert e.status == HTTPStatus.NOT_FOUND
            assert e.data["type"] == "https://api.backend.ai/probs/kernel-not-found"
            assert e.data["title"] == "Kernel Not Found"


async def test_fetch_invalid_method_async() -> None:
    async with AsyncSession():
        rqst = Request("STRANGE", "/")
        with pytest.raises(ValueError):
            async with rqst.fetch():
                pass


async def test_fetch_client_error_async(dummy_endpoint: str) -> None:
    with aioresponses() as m:
        async with AsyncSession():
            m.post(dummy_endpoint, exception=aiohttp.ClientConnectionError())
            rqst = Request("POST", "/")
            with pytest.raises(BackendClientError):
                async with rqst.fetch():
                    pass


@pytest.mark.xfail
async def test_fetch_cancellation_async(dummy_endpoint: str) -> None:
    # It seems that aiohttp swallows asyncio.CancelledError
    with aioresponses() as m:
        async with AsyncSession():
            m.post(dummy_endpoint, exception=asyncio.CancelledError())
            rqst = Request("POST", "/")
            with pytest.raises(asyncio.CancelledError):
                async with rqst.fetch():
                    pass


async def test_fetch_timeout_async(dummy_endpoint: str) -> None:
    with aioresponses() as m:
        async with AsyncSession():
            m.post(dummy_endpoint, exception=TimeoutError())
            rqst = Request("POST", "/")
            with pytest.raises(asyncio.TimeoutError):
                async with rqst.fetch():
                    pass


async def test_response_async(defconfig: APIConfig, dummy_endpoint: str) -> None:
    body = b'{"test": 5678}'
    with aioresponses() as m:
        m.post(
            dummy_endpoint + "function",
            status=HTTPStatus.OK,
            body=body,
            headers={"Content-Type": "application/json", "Content-Length": str(len(body))},
        )
        async with AsyncSession(config=defconfig):
            rqst = Request("POST", "/function")
            async with rqst.fetch() as resp:
                assert await resp.text() == '{"test": 5678}'
                assert await resp.json() == {"test": 5678}


async def test_fetch_preserves_date_header_override(dummy_endpoint: str) -> None:
    fixed_date = "Tue, 02 Sep 2025 08:00:00 GMT"
    with aioresponses() as m:
        m.post(dummy_endpoint + "function", status=HTTPStatus.OK, body=b"")
        async with AsyncSession():
            rqst = Request("POST", "function")
            async with rqst.fetch(headers={"Date": fixed_date}):
                pass

        sent_kwargs = next(iter(m.requests.values()))[0].kwargs
        assert sent_kwargs["headers"]["Date"] == fixed_date

    # The internal datetime must also reflect the override so that signing
    # done with `self.date` matches the Date header sent on the wire.
    assert rqst.date is not None
    assert rqst.date.year == 2025
    assert rqst.date.month == 9
    assert rqst.date.day == 2


async def test_fetch_default_date_header_when_no_override(dummy_endpoint: str) -> None:
    with aioresponses() as m:
        m.post(dummy_endpoint + "function", status=HTTPStatus.OK, body=b"")
        async with AsyncSession():
            rqst = Request("POST", "function")
            async with rqst.fetch():
                pass

        sent_kwargs = next(iter(m.requests.values()))[0].kwargs
        # Without override, fetch() auto-populates Date with self.date.isoformat().
        assert rqst.date is not None
        assert sent_kwargs["headers"]["Date"] == rqst.date.isoformat()


async def test_fetch_passes_through_arbitrary_header_overrides(dummy_endpoint: str) -> None:
    with aioresponses() as m:
        m.post(dummy_endpoint + "function", status=HTTPStatus.OK, body=b"")
        async with AsyncSession():
            rqst = Request("POST", "function")
            async with rqst.fetch(headers={"X-Custom-Header": "custom-value"}):
                pass

        sent_kwargs = next(iter(m.requests.values()))[0].kwargs
        assert sent_kwargs["headers"]["X-Custom-Header"] == "custom-value"