test_client.py

# License: MIT
# Copyright © 2024 Frequenz Energy-as-a-Service GmbH

"""Tests for the frequenz.client.dispatch package."""

import asyncio
import random
from dataclasses import replace
from datetime import timedelta
from functools import partial

import grpc
import pytest
from pytest import raises

from frequenz.client.common.microgrid import MicrogridId
from frequenz.client.dispatch.test.client import FakeClient, to_create_params
from frequenz.client.dispatch.test.fixtures import client, generator, sample
from frequenz.client.dispatch.test.generator import DispatchGenerator
from frequenz.client.dispatch.types import Dispatch, DispatchId, Event

# Ignore flake8 error in the rest of the file to use the same fixture names
# flake8: noqa[811]


def _update_metadata(dispatch: Dispatch, created: Dispatch) -> Dispatch:
    """Update the dispatch.

    Updates the id, create_time and update_time of the local dispatch
    with the values from the created dispatch.

    Args:
        dispatch: The local dispatch.
        created: The created dispatch.

    Returns:
        The updated dispatch.
    """
    return replace(
        dispatch,
        id=created.id,
        create_time=created.create_time,
        update_time=created.update_time,
        end_time=created.end_time,
    )


async def test_create_dispatch(client: FakeClient, sample: Dispatch) -> None:
    """Test creating a dispatch."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    dispatch = await client.create(**to_create_params(microgrid_id, sample))

    sample = _update_metadata(sample, dispatch)
    assert dispatch == sample


async def test_create_return_dispatch(
    client: FakeClient, generator: DispatchGenerator
) -> None:
    """Test creating a dispatch and returning the created dispatch."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    for _ in range(100):
        sample = generator.generate_dispatch()

        dispatch = await client.create(**to_create_params(microgrid_id, sample))

        sample = _update_metadata(sample, dispatch)

        assert dispatch == sample


async def test_create_duration_none(client: FakeClient, sample: Dispatch) -> None:
    """Test creating a dispatch with a None duration."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    sample = replace(sample, duration=None)
    dispatch = await client.create(**to_create_params(microgrid_id, sample))
    sample = _update_metadata(sample, dispatch)
    assert dispatch == sample


async def test_create_duration_0(client: FakeClient, sample: Dispatch) -> None:
    """Test creating a dispatch with a 0 duration."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    sample = replace(sample, duration=timedelta(minutes=0))
    dispatch = await client.create(**to_create_params(microgrid_id, sample))
    sample = _update_metadata(sample, dispatch)
    assert dispatch == sample


async def test_list_dispatches(
    client: FakeClient, generator: DispatchGenerator
) -> None:
    """Test listing dispatches."""
    microgrid_id = MicrogridId(random.randint(1, 100))

    client.set_dispatches(
        microgrid_id=microgrid_id,
        value=[generator.generate_dispatch() for _ in range(100)],
    )

    dispatches = client.list(microgrid_id=MicrogridId(1))
    async for page in dispatches:
        for dispatch in page:
            # First find matching id in client.dispatches, then compare
            service_side_dispatch = next(
                filter(
                    partial(lambda d, md: d.id == md.id, md=dispatch),
                    client.dispatches(microgrid_id=MicrogridId(1)),
                ),
                None,
            )

            assert dispatch == service_side_dispatch


# pylint: disable-next=too-many-locals
async def test_list_filter_queries(
    client: FakeClient, generator: DispatchGenerator
) -> None:
    """Test listing dispatches with filter queries."""
    microgrid_id = MicrogridId(1)

    all_dispatches = [generator.generate_dispatch() for _ in range(100)]
    client.set_dispatches(
        microgrid_id=microgrid_id,
        value=all_dispatches,
    )

    # Filter by type
    filter_type = all_dispatches[0].type
    dispatches = client.list(
        microgrid_id=microgrid_id,
        filter_queries=iter([filter_type]),
    )
    async for page in dispatches:
        for dispatch in page:
            assert dispatch.type == filter_type

    # Filter by id (needs to prefix with #)
    filter_target_id = all_dispatches[0].id
    dispatches = client.list(
        microgrid_id=microgrid_id,
        filter_queries=iter([f"#{int(filter_target_id)}"]),
    )
    async for page in dispatches:
        for dispatch in page:
            assert dispatch.id == filter_target_id

    # Mixed filter - validate OR behavior with type and id
    filter_mixed = [all_dispatches[3].type, f"#{int(all_dispatches[4].id)}"]
    dispatches = client.list(
        microgrid_id=microgrid_id,
        filter_queries=iter(filter_mixed),
    )
    async for page in dispatches:
        for dispatch in page:
            assert (
                dispatch.type == all_dispatches[3].type
                or dispatch.id == all_dispatches[4].id
            )

    # Test OR behavior with multiple types - validate both dispatches are found
    type1 = all_dispatches[5].type
    type2 = all_dispatches[6].type
    dispatches = client.list(
        microgrid_id=microgrid_id,
        filter_queries=iter([type1, type2]),
    )
    found_dispatches = []
    async for page in dispatches:
        for dispatch in page:
            assert dispatch.type in (type1, type2)
            found_dispatches.append(dispatch)

    # Verify both dispatches with the specified types are found
    assert len(found_dispatches) == 2
    assert any(dispatch.type == type1 for dispatch in found_dispatches)
    assert any(dispatch.type == type2 for dispatch in found_dispatches)

    # Test OR behavior with multiple IDs - validate both dispatches are found
    id1 = all_dispatches[7].id
    id2 = all_dispatches[8].id
    # Use numeric part only for ID queries (fake service expects just the number)
    dispatches = client.list(
        microgrid_id=microgrid_id,
        filter_queries=iter([f"#{int(id1)}", f"#{int(id2)}"]),
    )
    found_dispatches = []
    async for page in dispatches:
        for dispatch in page:
            assert dispatch.id in (id1, id2)
            found_dispatches.append(dispatch)

    # Verify both dispatches with the specified IDs are found
    assert len(found_dispatches) == 2
    assert any(dispatch.id == id1 for dispatch in found_dispatches)
    assert any(dispatch.id == id2 for dispatch in found_dispatches)

    # Test OR behavior with mixed types and IDs - validate all dispatches are found
    mixed_filter = [
        all_dispatches[9].type,
        f"#{int(all_dispatches[10].id)}",
        all_dispatches[11].type,
        f"#{int(all_dispatches[12].id)}",
    ]
    dispatches = client.list(
        microgrid_id=microgrid_id,
        filter_queries=iter(mixed_filter),
    )
    found_dispatches = []
    async for page in dispatches:
        for dispatch in page:
            assert (
                dispatch.type == all_dispatches[9].type
                or dispatch.id == all_dispatches[10].id
                or dispatch.type == all_dispatches[11].type
                or dispatch.id == all_dispatches[12].id
            )
            found_dispatches.append(dispatch)

    # Verify all four dispatches with the specified types/IDs are found
    assert len(found_dispatches) == 4
    assert any(dispatch.type == all_dispatches[9].type for dispatch in found_dispatches)
    assert any(dispatch.id == all_dispatches[10].id for dispatch in found_dispatches)
    assert any(
        dispatch.type == all_dispatches[11].type for dispatch in found_dispatches
    )
    assert any(dispatch.id == all_dispatches[12].id for dispatch in found_dispatches)


async def test_list_dispatches_no_duration(
    client: FakeClient, generator: DispatchGenerator
) -> None:
    """Test listing dispatches with a None duration."""
    microgrid_id = MicrogridId(random.randint(1, 100))

    client.set_dispatches(
        microgrid_id=microgrid_id,
        value=[
            replace(generator.generate_dispatch(), duration=None) for _ in range(100)
        ],
    )

    dispatches = client.list(microgrid_id=MicrogridId(1))
    async for page in dispatches:
        for dispatch in page:
            assert dispatch in client.dispatches(microgrid_id=MicrogridId(1))


async def test_list_create_dispatches(
    client: FakeClient, generator: DispatchGenerator
) -> None:
    """Test listing dispatches."""
    microgrid_id = MicrogridId(random.randint(1, 100))

    # Test with empty list
    page = await anext(client.list(microgrid_id=microgrid_id))
    assert not any(page)

    for _ in range(100):
        sample = generator.generate_dispatch()
        await client.create(
            **to_create_params(microgrid_id=microgrid_id, dispatch=sample)
        )

        dispatch = None
        async for page in client.list(microgrid_id=microgrid_id):
            for dispatch in page:
                pass

        if dispatch is None:
            raise AssertionError("Dispatch not found")

        sample = _update_metadata(sample, dispatch)
        assert dispatch == sample


async def test_update_dispatch(client: FakeClient, sample: Dispatch) -> None:
    """Test updating a dispatch."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    dispatch = await client.create(**to_create_params(microgrid_id, sample))

    sample = _update_metadata(sample, dispatch)
    assert dispatch == sample

    await client.update(
        microgrid_id=microgrid_id,
        dispatch_id=dispatch.id,
        new_fields={"recurrence.interval": 4},
    )
    assert client.dispatches(microgrid_id)[0].recurrence.interval == 4


async def test_update_dispatch_to_no_duration(
    client: FakeClient, sample: Dispatch
) -> None:
    """Test updating the duration field of a dispatch to None."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    client.set_dispatches(
        microgrid_id=microgrid_id,
        value=[replace(sample, duration=timedelta(minutes=10))],
    )

    await client.update(
        microgrid_id=microgrid_id,
        dispatch_id=sample.id,
        new_fields={"duration": None},
    )
    assert client.dispatches(microgrid_id)[0].duration is None


async def test_update_dispatch_to_0_duration(
    client: FakeClient, sample: Dispatch
) -> None:
    """Test updating the duration field of a dispatch to 0."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    client.set_dispatches(
        microgrid_id=microgrid_id,
        value=[replace(sample, duration=timedelta(minutes=10))],
    )

    await client.update(
        microgrid_id=microgrid_id,
        dispatch_id=sample.id,
        new_fields={"duration": timedelta(minutes=0)},
    )
    assert client.dispatches(microgrid_id)[0].duration == timedelta(minutes=0)


async def test_update_dispatch_from_no_duration(
    client: FakeClient, sample: Dispatch
) -> None:
    """Test updating the duration field of a dispatch from None."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    client.set_dispatches(
        microgrid_id=microgrid_id, value=[replace(sample, duration=None)]
    )

    await client.update(
        microgrid_id=microgrid_id,
        dispatch_id=sample.id,
        new_fields={"duration": timedelta(minutes=10)},
    )
    assert client.dispatches(microgrid_id)[0].duration == timedelta(minutes=10)


async def test_update_dispatch_fail(client: FakeClient, sample: Dispatch) -> None:
    """Test updating the type and dry_run fields of a dispatch."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    dispatch = await client.create(**to_create_params(microgrid_id, sample))

    assert dispatch is not None

    sample = _update_metadata(sample, dispatch)
    assert dispatch == sample

    for field, value in [
        ("type", "new_type"),
        ("dry_run", True),
        ("is_dry_run", True),
    ]:
        with raises(ValueError):
            await client.update(
                microgrid_id=microgrid_id,
                dispatch_id=dispatch.id,
                new_fields={field: value},
            )


async def test_get_dispatch(client: FakeClient, sample: Dispatch) -> None:
    """Test getting a dispatch."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    dispatch = await client.create(**to_create_params(microgrid_id, sample))

    sample = _update_metadata(sample, dispatch)
    assert dispatch == sample

    assert (
        await client.get(microgrid_id=microgrid_id, dispatch_id=dispatch.id) == dispatch
    )


async def test_get_dispatch_no_duration(client: FakeClient, sample: Dispatch) -> None:
    """Test getting a dispatch with a None duration."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    sample = replace(sample, duration=None)
    dispatch = await client.create(**to_create_params(microgrid_id, sample))

    sample = _update_metadata(sample, dispatch)
    assert dispatch == sample

    assert (
        await client.get(microgrid_id=microgrid_id, dispatch_id=dispatch.id) == dispatch
    )


async def test_get_dispatch_fail(client: FakeClient) -> None:
    """Test getting a non-existent dispatch."""
    with raises(grpc.RpcError):
        await client.get(microgrid_id=MicrogridId(1), dispatch_id=DispatchId(1))


async def test_delete_dispatch(client: FakeClient, sample: Dispatch) -> None:
    """Test deleting a dispatch."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    dispatch = await client.create(**to_create_params(microgrid_id, sample))

    sample = _update_metadata(sample, dispatch)
    assert dispatch == sample

    await client.delete(microgrid_id=microgrid_id, dispatch_id=dispatch.id)

    assert len(client.dispatches(microgrid_id)) == 0


async def test_delete_dispatch_fail(client: FakeClient) -> None:
    """Test deleting a non-existent dispatch."""
    with raises(grpc.RpcError):
        await client.delete(microgrid_id=MicrogridId(1), dispatch_id=DispatchId(1))


@pytest.mark.parametrize("call_twice", [True, False])
async def test_dispatch_stream(
    client: FakeClient, sample: Dispatch, call_twice: bool
) -> None:
    """Test dispatching a stream of dispatches."""
    microgrid_id = MicrogridId(random.randint(1, 100))
    dispatches = [sample, sample, sample]

    stream = client.stream(microgrid_id)

    if call_twice:
        # Call function again to test behavior if stream already exists
        _ = client.stream(microgrid_id)

    async def expect(dispatch: Dispatch, event: Event) -> None:
        message = await stream.receive()
        assert message.dispatch == dispatch
        assert message.event == event

    # Give stream some time to start
    await asyncio.sleep(0.1)

    # Add a new dispatch
    dispatches[0] = await client.create(**to_create_params(microgrid_id, dispatches[0]))
    # Expect the first dispatch event
    await expect(dispatches[0], Event.CREATED)

    # Add a new dispatch
    dispatches[1] = await client.create(**to_create_params(microgrid_id, dispatches[1]))
    # Expect the second dispatch
    await expect(dispatches[1], Event.CREATED)

    # Add a new dispatch
    dispatches[2] = await client.create(**to_create_params(microgrid_id, dispatches[2]))
    # Expect the third dispatch
    await expect(dispatches[2], Event.CREATED)

    # Update the first dispatch
    dispatches[0] = await client.update(
        microgrid_id=microgrid_id,
        dispatch_id=dispatches[0].id,
        new_fields={"start_time": dispatches[0].start_time + timedelta(minutes=1)},
    )

    # Expect the first dispatch update
    await expect(dispatches[0], Event.UPDATED)

    # Delete the first dispatch
    await client.delete(microgrid_id=microgrid_id, dispatch_id=dispatches[0].id)

    # Expect the first dispatch deletion
    await expect(dispatches[0], Event.DELETED)