EOS/tests/test_pydantic.py

from typing import Optional

import pandas as pd
import pendulum
import pytest
from pydantic import Field, ValidationError

from akkudoktoreos.core.pydantic import (
    PydanticBaseModel,
    PydanticDateTimeData,
    PydanticDateTimeDataFrame,
    PydanticDateTimeSeries,
    PydanticModelNestedValueMixin,
)
from akkudoktoreos.utils.datetimeutil import compare_datetimes, to_datetime


class PydanticTestModel(PydanticBaseModel):
    datetime_field: pendulum.DateTime = Field(
        ..., description="A datetime field with pendulum support."
    )
    optional_field: Optional[str] = Field(default=None, description="An optional field.")


class Address(PydanticBaseModel):
    city: Optional[str] = None
    postal_code: Optional[str] = None


class User(PydanticBaseModel):
    name: str
    addresses: Optional[list[Address]] = None
    settings: Optional[dict[str, str]] = None


class TestPydanticModelNestedValueMixin:
    """Umbrella test class to group all test cases for `PydanticModelNestedValueMixin`."""

    @pytest.fixture
    def user_instance(self):
        """Fixture to initialize a sample User instance."""
        return User(name="Alice", addresses=None, settings=None)

    def test_get_key_types_for_simple_field(self):
        """Test _get_key_types for a simple string field."""
        key_types = PydanticModelNestedValueMixin._get_key_types(User, "name")
        assert key_types == [str], f"Expected [str], got {key_types}"

    def test_get_key_types_for_list_of_models(self):
        """Test _get_key_types for a list of Address models."""
        key_types = PydanticModelNestedValueMixin._get_key_types(User, "addresses")
        assert key_types == [list, Address], f"Expected [list, Address], got {key_types}"

    def test_get_key_types_for_dict_field(self):
        """Test _get_key_types for a dictionary field."""
        key_types = PydanticModelNestedValueMixin._get_key_types(User, "settings")
        assert key_types == [dict, str], f"Expected [dict, str], got {key_types}"

    def test_get_key_types_for_optional_field(self):
        """Test _get_key_types correctly handles Optional fields."""
        key_types = PydanticModelNestedValueMixin._get_key_types(Address, "city")
        assert key_types == [str], f"Expected [str], got {key_types}"

    def test_get_key_types_for_non_existent_field(self):
        """Test _get_key_types raises an error for non-existent field."""
        with pytest.raises(TypeError):
            PydanticModelNestedValueMixin._get_key_types(User, "unknown_field")

    def test_set_nested_value_in_model(self, user_instance):
        """Test setting nested value in a model field (Address -> city)."""
        assert user_instance.addresses is None

        user_instance.set_nested_value("addresses/0/city", "New York")

        assert user_instance.addresses is not None
        assert user_instance.addresses[0].city == "New York", "The city should be set to 'New York'"

    def test_set_nested_value_in_dict(self, user_instance):
        """Test setting nested value in a dictionary field (settings -> theme)."""
        assert user_instance.settings is None

        user_instance.set_nested_value("settings/theme", "dark")

        assert user_instance.settings is not None
        assert user_instance.settings["theme"] == "dark", "The theme should be set to 'dark'"

    def test_set_nested_value_in_list(self, user_instance):
        """Test setting nested value in a list of models (addresses -> 1 -> city)."""
        user_instance.set_nested_value("addresses/1/city", "Los Angeles")

        # Check if the city in the second address is set correctly
        assert user_instance.addresses[1].city == "Los Angeles", (
            "The city at index 1 should be set to 'Los Angeles'"
        )

    def test_set_nested_value_in_optional_field(self, user_instance):
        """Test setting value in an Optional field (addresses)."""
        user_instance.set_nested_value("addresses/0", Address(city="Chicago"))

        # Check if the first address is set correctly
        assert user_instance.addresses is not None
        assert user_instance.addresses[0].city == "Chicago", "The city should be set to 'Chicago'"

    def test_set_nested_value_with_empty_list(self):
        """Test setting value in an empty list of models."""
        user = User(name="Bob", addresses=[])
        user.set_nested_value("addresses/0/city", "Seattle")

        assert user.addresses is not None
        assert user.addresses[0].city == "Seattle", (
            "The first address should have the city 'Seattle'"
        )

    def test_set_nested_value_with_missing_key_in_dict(self, user_instance):
        """Test setting value in a dict when the key does not exist."""
        user_instance.set_nested_value("settings/language", "English")

        assert user_instance.settings["language"] == "English", (
            "The language setting should be 'English'"
        )

    def test_set_nested_value_for_non_existent_field(self):
        """Test attempting to set value for a non-existent field."""
        user = User(name="John")

        with pytest.raises(ValueError):
            user.set_nested_value("non_existent_field", "Some Value")

    def test_set_nested_value_with_invalid_type(self, user_instance):
        """Test setting value with an invalid type."""
        with pytest.raises(ValueError):
            user_instance.set_nested_value(
                "addresses/0/city", 1234
            )  # city should be a string, not an integer

    def test_set_nested_value_with_model_initialization(self):
        """Test setting a value in a model that should initialize a missing model."""
        user = User(name="James", addresses=None)
        user.set_nested_value("addresses/0/city", "Boston")

        assert user.addresses is not None
        assert user.addresses[0].city == "Boston", "The city should be set to 'Boston'"
        assert isinstance(user.addresses[0], Address), (
            "The first address should be an instance of Address"
        )


class TestPydanticBaseModel:
    def test_valid_pendulum_datetime(self):
        dt = pendulum.now()
        model = PydanticTestModel(datetime_field=dt)
        assert model.datetime_field == dt

    def test_invalid_datetime_string(self):
        with pytest.raises(ValidationError, match="Input should be an instance of DateTime"):
            PydanticTestModel(datetime_field="invalid_datetime")

    def test_iso8601_serialization(self):
        dt = pendulum.datetime(2024, 12, 21, 15, 0, 0)
        model = PydanticTestModel(datetime_field=dt)
        serialized = model.to_dict()
        expected_dt = to_datetime(dt)
        result_dt = to_datetime(serialized["datetime_field"])
        assert compare_datetimes(result_dt, expected_dt)

    def test_reset_to_defaults(self):
        dt = pendulum.now()
        model = PydanticTestModel(datetime_field=dt, optional_field="some value")
        model.reset_to_defaults()
        assert model.datetime_field == dt
        assert model.optional_field is None

    def test_from_dict_and_to_dict(self):
        dt = pendulum.now()
        model = PydanticTestModel(datetime_field=dt)
        data = model.to_dict()
        restored_model = PydanticTestModel.from_dict(data)
        assert restored_model.datetime_field == dt

    def test_to_json_and_from_json(self):
        dt = pendulum.now()
        model = PydanticTestModel(datetime_field=dt)
        json_data = model.to_json()
        restored_model = PydanticTestModel.from_json(json_data)
        assert restored_model.datetime_field == dt


class TestPydanticDateTimeData:
    def test_valid_list_lengths(self):
        data = {
            "timestamps": ["2024-12-21T15:00:00+00:00"],
            "values": [100],
        }
        model = PydanticDateTimeData(root=data)
        assert pendulum.parse(model.root["timestamps"][0]) == pendulum.parse(
            "2024-12-21T15:00:00+00:00"
        )

    def test_invalid_list_lengths(self):
        data = {
            "timestamps": ["2024-12-21T15:00:00+00:00"],
            "values": [100, 200],
        }
        with pytest.raises(
            ValidationError, match="All lists in the dictionary must have the same length"
        ):
            PydanticDateTimeData(root=data)


class TestPydanticDateTimeDataFrame:
    def test_valid_dataframe(self):
        df = pd.DataFrame(
            {
                "value": [100, 200],
            },
            index=pd.to_datetime(["2024-12-21", "2024-12-22"]),
        )
        model = PydanticDateTimeDataFrame.from_dataframe(df)
        result = model.to_dataframe()

        # Check index
        assert len(result.index) == len(df.index)
        for i, dt in enumerate(df.index):
            expected_dt = to_datetime(dt)
            result_dt = to_datetime(result.index[i])
            assert compare_datetimes(result_dt, expected_dt).equal


class TestPydanticDateTimeSeries:
    def test_valid_series(self):
        series = pd.Series([100, 200], index=pd.to_datetime(["2024-12-21", "2024-12-22"]))
        model = PydanticDateTimeSeries.from_series(series)
        result = model.to_series()

        # Check index
        assert len(result.index) == len(series.index)
        for i, dt in enumerate(series.index):
            expected_dt = to_datetime(dt)
            result_dt = to_datetime(result.index[i])
            assert compare_datetimes(result_dt, expected_dt).equal