Use new type DispatchEvent for new/updated/deleted dispatches

Also add some tests for it.

Signed-off-by: Mathias L. Baumann <[email protected]>

Author: Marenz

README.md

@@ -15,22 +15,30 @@ It provides two channels for clients:
 ## Example Usage
 
 ```python
-async def run():
-    # dispatch helper sends out dispatches when they are due
-    dispatch_arrived = dispatch_helper.new_dispatches().new_receiver()
-    dispatch_ready = dispatch_helper.ready_dispatches().new_receiver()
-
-    async for selected in select(dispatch_ready, dispatch_arrived):
-        if selected_from(selected, dispatch_arrived):
-            dispatch = selected.value
-            match dispatch.type:
-                case DISPATCH_TYPE_BATTERY_CHARGE:
-                    battery_pool = microgrid.battery_pool(dispatch.battery_set, task_id)
-                    battery_pool.set_charge(dispatch.power)
-            ...
-        if selected_from(selected, dispatch_ready):
-            dispatch = selected.value
-            log.info("New dispatch arrived %s", dispatch)
+    async def run():
+        # dispatch helper sends out dispatches when they are due
+        dispatch_arrived = dispatch_helper.updated_dispatches().new_receiver()
+        dispatch_ready = dispatch_helper.ready_dispatches().new_receiver()
+
+        async for selected in select(dispatch_ready, dispatch_arrived):
+            if selected_from(selected, dispatch_ready):
+                dispatch = selected.value
+                match dispatch.type:
+                    case DISPATCH_TYPE_BATTERY_CHARGE:
+                        battery_pool = microgrid.battery_pool(dispatch.battery_set, task_id)
+                        battery_pool.set_charge(dispatch.power)
+                ...
+            if selected_from(selected, dispatch_arrived):
+                match selected.value:
+                    case Created(dispatch):
+                        log.info("New dispatch arrived %s", dispatch)
+                        ...
+                    case Updated(dispatch):
+                        log.info("Dispatch updated %s", dispatch)
+                        ...
+                    case Deleted(dispatch):
+                        log.info("Dispatch deleted %s", dispatch)
+                        ...
 ```
 
 ## Supported Platforms

src/frequenz/dispatch/__init__.py

@@ -7,7 +7,7 @@
 from frequenz.channels import Broadcast, Receiver
 from frequenz.client.dispatch.types import Dispatch
 
-from frequenz.dispatch.actor import DispatchActor
+from frequenz.dispatch.actor import DispatchActor, DispatchEvent
 
 __all__ = ["Dispatcher"]
 
@@ -41,26 +41,26 @@ def __init__(
             svc_addr: The service address.
         """
         self._ready_channel = Broadcast[Dispatch]("ready_dispatches")
-        self._new_channel = Broadcast[Dispatch]("new_dispatches")
+        self._updated_channel = Broadcast[DispatchEvent]("new_dispatches")
         self._actor = DispatchActor(
             microgrid_id,
             grpc_channel,
             svc_addr,
-            self._new_channel.new_sender(),
+            self._updated_channel.new_sender(),
             self._ready_channel.new_sender(),
         )
 
     async def start(self) -> None:
         """Start the actor."""
         self._actor.start()
 
-    def new_dispatches(self) -> Receiver[Dispatch]:
-        """Return new dispatches receiver.
+    def updated_dispatches(self) -> Receiver[DispatchEvent]:
+        """Return new, updated or deleted dispatches receiver.
 
         Returns:
             A new receiver for new dispatches.
         """
-        return self._new_channel.new_receiver()
+        return self._updated_channel.new_receiver()
 
     def ready_dispatches(self) -> Receiver[Dispatch]:
         """Return ready dispatches receiver.

src/frequenz/dispatch/actor.py

@@ -5,6 +5,7 @@
 
 import asyncio
 import logging
+from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
 from typing import cast
 
@@ -51,6 +52,36 @@
 """To map from our Weekday enum to the dateutil library enum."""
 
 
+@dataclass(frozen=True)
+class _DispatchEventBase:
+    dispatch: Dispatch
+    """The dispatch that this event is about.
+
+    Objects of this base class are sent over the channel when a dispatch is
+    created, updated or deleted.
+    """
+
+
+class Created(_DispatchEventBase):
+    """Wraps a dispatch that was created."""
+
+
+class Updated(_DispatchEventBase):
+    """Wraps a dispatch that was updated."""
+
+
+class Deleted(_DispatchEventBase):
+    """Wraps a dispatch that was deleted."""
+
+
+DispatchEvent = Created | Updated | Deleted
+"""Type that is sent over the channel for dispatch updates.
+
+This type is used to send dispatches that were created, updated or deleted
+over the channel.
+"""
+
+
 class DispatchActor(Actor):
     """Dispatch actor.
 
@@ -66,7 +97,7 @@ def __init__(
         microgrid_id: int,
         grpc_channel: grpc.aio.Channel,
         svc_addr: str,
-        updated_dispatch_sender: Sender[Dispatch],
+        updated_dispatch_sender: Sender[DispatchEvent],
         ready_dispatch_sender: Sender[Dispatch],
     ) -> None:
         """Initialize the actor.
@@ -112,11 +143,11 @@ async def _fetch(self) -> None:
                 if not old_dispatch:
                     self._update_dispatch_schedule(dispatch, None)
                     _logger.info("New dispatch: %s", dispatch)
-                    await self._updated_dispatch_sender.send(dispatch)
+                    await self._updated_dispatch_sender.send(Created(dispatch=dispatch))
                 elif dispatch.update_time != old_dispatch.update_time:
                     self._update_dispatch_schedule(dispatch, old_dispatch)
                     _logger.info("Updated dispatch: %s", dispatch)
-                    await self._updated_dispatch_sender.send(dispatch)
+                    await self._updated_dispatch_sender.send(Updated(dispatch=dispatch))
 
         except grpc.aio.AioRpcError as error:
             _logger.error("Error fetching dispatches: %s", error)
@@ -125,6 +156,7 @@ async def _fetch(self) -> None:
 
         for dispatch in old_dispatches.values():
             _logger.info("Deleted dispatch: %s", dispatch)
+            await self._updated_dispatch_sender.send(Deleted(dispatch=dispatch))
             if task := self._scheduled.pop(dispatch.id, None):
                 task.cancel()
 

tests/test_frequenz_dispatch.py

@@ -4,10 +4,10 @@
 """Tests for the frequenz.dispatch.actor package."""
 
 import asyncio
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from datetime import datetime, timedelta, timezone
 from random import randint
-from typing import AsyncIterator, Iterator
+from typing import AsyncIterator, Iterator, TypeVar
 from unittest.mock import MagicMock
 
 import async_solipsism
@@ -16,10 +16,16 @@
 from frequenz.channels._broadcast import Sender
 from frequenz.client.dispatch.test.client import FakeClient, to_create_params
 from frequenz.client.dispatch.test.generator import DispatchGenerator
-from frequenz.client.dispatch.types import Dispatch
+from frequenz.client.dispatch.types import Dispatch, Frequency
 from pytest import fixture
 
-from frequenz.dispatch.actor import DispatchActor
+from frequenz.dispatch.actor import (
+    Created,
+    Deleted,
+    DispatchActor,
+    DispatchEvent,
+    Updated,
+)
 
 
 # This method replaces the event loop for all tests in the file.
@@ -45,13 +51,13 @@ def _now() -> datetime:
     return datetime.now(tz=timezone.utc)
 
 
-@dataclass
+@dataclass(frozen=True)
 class ActorTestEnv:
     """Test environment for the actor."""
 
     actor: DispatchActor
     """The actor under test."""
-    updated_dispatches: Receiver[Dispatch]
+    updated_dispatches: Receiver[DispatchEvent]
     """The receiver for updated dispatches."""
     ready_dispatches: Receiver[Dispatch]
     """The receiver for ready dispatches."""
@@ -64,8 +70,9 @@ class ActorTestEnv:
 @fixture
 async def actor_env() -> AsyncIterator[ActorTestEnv]:
     """Return an actor test environment."""
+    T = TypeVar("T")
 
-    class YieldingSender(Sender[Dispatch]):
+    class YieldingSender(Sender[T]):
         """A sender that yields after sending.
 
         For testing we want to manipulate the time after a call to send.
@@ -74,15 +81,12 @@ class YieldingSender(Sender[Dispatch]):
         opportunity to manipulate the time.
         """
 
-        def __init__(self, channel: Broadcast[Dispatch]) -> None:
-            super().__init__(channel)
-
-        async def send(self, msg: Dispatch) -> None:
+        async def send(self, msg: T) -> None:
             """Send the value and yield."""
             await super().send(msg)
             await asyncio.sleep(1)
 
-    updated_dispatches = Broadcast[Dispatch]("updated_dispatches")
+    updated_dispatches = Broadcast[DispatchEvent]("updated_dispatches")
     ready_dispatches = Broadcast[Dispatch]("ready_dispatches")
     microgrid_id = randint(1, 100)
 
@@ -122,13 +126,96 @@ async def test_new_dispatch_created(
 ) -> None:
     """Test that a new dispatch is created."""
     sample = generator.generate_dispatch(actor_env.microgrid_id)
+
+    await _test_new_dispatch_created(actor_env, sample)
+
+
+async def _test_new_dispatch_created(
+    actor_env: ActorTestEnv,
+    sample: Dispatch,
+) -> Dispatch:
+    """Test that a new dispatch is created.
+
+    Args:
+        actor_env: The actor environment
+        sample: The sample dispatch to create
+
+    Returns:
+        The sample dispatch that was created
+    """
     await actor_env.client.create(**to_create_params(sample))
 
-    dispatch = await actor_env.updated_dispatches.receive()
-    sample.update_time = dispatch.update_time
-    sample.create_time = dispatch.create_time
-    sample.id = dispatch.id
-    assert dispatch == sample
+    dispatch_event = await actor_env.updated_dispatches.receive()
+
+    match dispatch_event:
+        case Deleted(dispatch) | Updated(dispatch):
+            assert False, "Expected a created event"
+        case Created(dispatch):
+            sample.update_time = dispatch.update_time
+            sample.create_time = dispatch.create_time
+            sample.id = dispatch.id
+            assert dispatch == sample
+
+    return sample
+
+
+async def test_existing_dispatch_updated(
+    actor_env: ActorTestEnv,
+    generator: DispatchGenerator,
+    fake_time: time_machine.Coordinates,
+) -> None:
+    """Test that an existing dispatch is updated."""
+    sample = generator.generate_dispatch(actor_env.microgrid_id)
+    sample.active = False
+    sample.recurrence.frequency = Frequency.DAILY
+
+    fake_time.shift(timedelta(seconds=1))
+
+    await _test_new_dispatch_created(actor_env, sample)
+    fake_time.shift(timedelta(seconds=1))
+
+    await actor_env.client.update(
+        sample.id,
+        new_fields={
+            "active": True,
+            "recurrence.frequency": Frequency.UNSPECIFIED,
+        },
+    )
+    fake_time.shift(timedelta(seconds=1))
+
+    dispatch_event = await actor_env.updated_dispatches.receive()
+    match dispatch_event:
+        case Created(dispatch) | Deleted(dispatch):
+            assert False, "Expected an updated event"
+        case Updated(dispatch):
+            sample.update_time = dispatch.update_time
+            sample.active = True
+            sample.recurrence = replace(
+                sample.recurrence, frequency=Frequency.UNSPECIFIED
+            )
+            assert dispatch == sample
+
+
+async def test_existing_dispatch_deleted(
+    actor_env: ActorTestEnv,
+    generator: DispatchGenerator,
+    fake_time: time_machine.Coordinates,
+) -> None:
+    """Test that an existing dispatch is deleted."""
+    sample = generator.generate_dispatch(actor_env.microgrid_id)
+
+    await _test_new_dispatch_created(actor_env, sample)
+
+    await actor_env.client.delete(sample.id)
+    fake_time.shift(timedelta(seconds=1))
+
+    dispatch_event = await actor_env.updated_dispatches.receive()
+    match dispatch_event:
+        case Created(dispatch) | Updated(dispatch):
+            assert False, "Expected a deleted event"
+        case Deleted(dispatch):
+            sample.create_time = dispatch.create_time
+            assert dispatch == sample
 
 
 async def test_dispatch_schedule(