_test_utils.py

"""Define testing utils that need to imported."""

# pragma: exclude file

from __future__ import annotations

from enum import Enum

try:
    import pytest
except ImportError as exc:
    raise ImportError(
        "You must `pip install 'mp-api[test]' to use these testing utilities."
    ) from exc

import os
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Callable, Sequence
    from typing import Any

requires_api_key = pytest.mark.skipif(
    os.getenv("MP_API_KEY") is None,
    reason="No API key found.",
)

NUM_DOCS = 5


def client_search_testing(
    search_method: Callable,
    excluded_params: list[str],
    alt_name_dict: dict[str, str],
    custom_field_tests: dict[str, Any],
    sub_doc_fields: list[str],
    int_bounds: tuple[int, int] = (-100, 100),
    float_bounds: tuple[float, float] = (-100.12, 100.12),
):
    """Function to test a client using its search method.
    Each parameter is used to query for data, which is then checked.

    Args:
        search_method (Callable): Client search method
        excluded_params (list[str]): List of parameters to exclude from testing
        alt_name_dict (dict[str, str]): Alternative names for parameters used in the projection and subsequent data checking
        custom_field_tests (dict[str, Any]): Custom queries for specific fields.
        sub_doc_fields (list[str]): Prefixes for fields to check in resulting data. Useful when data to be tested is nested.
        int_bounds (tuple[int,int]) : integer bounds to use in testing int-type query arguments
        float_bounds (tuple[float,float]) : float bounds to use in testing float-type query arguments
    """
    if search_method is None:
        return
    # Get list of parameters
    param_tuples = list(search_method.__annotations__.items())

    # Query API for each numeric and boolean parameter and check if returned
    for entry in param_tuples:
        param = entry[0]

        if param not in excluded_params + ["return"]:
            param_type = entry[1]
            q: dict[str, Any] = {"chunk_size": 1, "num_chunks": 1}

            if "tuple[int, int]" in param_type:
                q[param] = int_bounds
            elif "tuple[float, float]" in param_type:
                q[param] = float_bounds
            elif "bool" in param_type:
                q[param] = False
            elif param in custom_field_tests:
                q[param] = custom_field_tests[param]
            else:
                raise ValueError(
                    f"Parameter '{param}' with type '{param_type}' was not "
                    "properly identified in the generic search method test."
                )

            if len(docs := search_method(**q)) > 0:
                doc = docs[0].model_dump()
            else:
                raise ValueError("No documents returned")

            for sub_field in sub_doc_fields:
                if sub_field in doc:
                    doc = doc[sub_field]

            assert doc[alt_name_dict.get(param, param)] is not None


def client_pagination(
    search_method: Callable, id_name: str, additional_fields: list[str] | None = None
) -> None:
    """Test pagination on an endpoint.

    Args:
    search_method (Callable) : Client search method to use
    id_name (str) : the name of a field which uniquely indexes a series of documents
    additional_fields (list of str) : Optional other fields to retrieve.

    Raises:
    AssertionError if pagination does not result in unique sets of documents
    """
    fields = [id_name, *(additional_fields or [])]
    page_1 = search_method(_page=1, chunk_size=NUM_DOCS, fields=fields)
    page_2 = search_method(_page=2, chunk_size=NUM_DOCS, fields=fields)
    assert all(len(results) == NUM_DOCS for results in (page_1, page_2))
    assert {str(getattr(doc, id_name)) for doc in page_1}.intersection(
        {str(getattr(doc, id_name)) for doc in page_2}
    ) == set()


def client_sort(
    search_method: Callable,
    sort_fields: str | Sequence[str],
    aux_query: dict[str, Any] | None = None,
    default_fields: tuple[str, ...] = ("deprecated", "material_id"),
):
    """Test sorting on an endpoint.

    Args:
    search_method (Callable) : Client search method to use
    sort_fields (str or Sequence of str) : fields to sort on
    aux_query (dict) : auxiliary query needed to filter documents
    default_fields (list): default fields to return

    Raises:
    AssertionError if sorting in ascending or descending order does not work.
    """

    def _normalize(doc, field: str):
        v = getattr(doc, field)
        # serialize enums
        return v.value if isinstance(v, Enum) else v

    user_query = {
        k: v
        for k, v in (aux_query or {}).items()
        if k not in ("_page", "_sort_fields", "chunk_size", "fields")
    }
    for sort_field in [sort_fields] if isinstance(sort_fields, str) else sort_fields:

        asc = search_method(
            _page=1,
            _sort_fields=sort_field,
            chunk_size=NUM_DOCS,
            fields=[sort_field, *default_fields],
            **user_query,
        )
        desc = search_method(
            _page=1,
            _sort_fields=f"-{sort_field}",
            chunk_size=NUM_DOCS,
            fields=[sort_field],
        )

        idxs = list(range(NUM_DOCS))
        assert sorted(idxs, key=lambda idx: _normalize(asc[idx], sort_field)) == idxs

        assert (
            sorted(
                idxs,
                key=lambda idx: _normalize(desc[idx], sort_field),
                reverse=True,
            )
            == idxs
        )