Remove streaming interface of StateTracker

... because it is no longer getting exposed to users directly,  but
instead through a hybrid method that will be introduced in the next
commit.

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/frequenz/sdk/timeseries/ev_charger_pool/_state_tracker.py

@@ -6,12 +6,10 @@
 from __future__ import annotations
 
 import asyncio
-from collections.abc import Iterator
-from dataclasses import dataclass
 from enum import Enum
 from typing import Optional
 
-from frequenz.channels import Broadcast, Receiver
+from frequenz.channels import Receiver
 from frequenz.channels.util import Merge
 
 from frequenz.sdk import microgrid
@@ -62,44 +60,6 @@ def from_ev_charger_data(cls, data: EVChargerData) -> EVChargerState:
         return EVChargerState.IDLE
 
 
-@dataclass(frozen=True)
-class EVChargerPoolStates:
-    """States of all EV Chargers in the pool."""
-
-    _states: dict[int, EVChargerState]
-    _changed_component: Optional[int] = None
-
-    def __iter__(self) -> Iterator[tuple[int, EVChargerState]]:
-        """Iterate over states of all EV Chargers.
-
-        Returns:
-            An iterator over all EV Charger states.
-        """
-        return iter(self._states.items())
-
-    def latest_change(self) -> Optional[tuple[int, EVChargerState]]:
-        """Return the most recent EV Charger state change.
-
-        The first `EVChargerPoolStates` instance created by a `StateTracker` will just
-        be a representation of the states of all EV Chargers.  At that point, the most
-        recent change in state of an ev charger will be unknown, so this function will
-        return `None`.
-
-        Returns:
-            None, when the most recent change is unknown.  Otherwise, a tuple with
-                the component ID of an EV Charger that just had a state change, and its
-                new state.
-        """
-        if self._changed_component is None:
-            return None
-        return (
-            self._changed_component,
-            self._states.setdefault(
-                self._changed_component, EVChargerState.UNSPECIFIED
-            ),
-        )
-
-
 class StateTracker:
     """A class for keeping track of the states of all EV Chargers in a pool."""
 
@@ -110,74 +70,50 @@ def __init__(self, component_ids: set[int]) -> None:
             component_ids: EV Charger component ids to track the states of.
         """
         self._component_ids = component_ids
-        self._channel = Broadcast[EVChargerPoolStates](
-            "EVCharger States", resend_latest=True
-        )
-        self._task: Optional[asyncio.Task[None]] = None
+        self._task: asyncio.Task[None] = asyncio.create_task(self._run())
         self._merged_stream: Optional[Merge[EVChargerData]] = None
-        self._states: dict[int, EVChargerState] = {}
 
-    def _get(self) -> EVChargerPoolStates:
-        """Get a representation of the current states of all EV Chargers.
+        # Initialize all components to the `MISSING` state.  This will change as data
+        # starts arriving from the individual components.
+        self._states: dict[int, EVChargerState] = {
+            component_id: EVChargerState.MISSING for component_id in component_ids
+        }
+
+    def get(self, component_id: int) -> EVChargerState:
+        """Return the current state of the EV Charger with the given component ID.
+
+        Args:
+            component_id: id of the EV Charger whose state is being fetched.
 
         Returns:
-            An `EVChargerPoolStates` instance.
+            An `EVChargerState` value corresponding to the given component id.
         """
-        return EVChargerPoolStates(self._states)
+        return self._states[component_id]
 
     def _update(
         self,
         data: EVChargerData,
-    ) -> Optional[EVChargerPoolStates]:
+    ) -> None:
         """Update the state of an EV Charger, from a new data point.
 
         Args:
             data: component data from the microgrid, for an EV Charger in the pool.
-
-        Returns:
-            A new `EVChargerPoolStates` instance representing all the EV Chargers in
-                the pool, in case there has been a state change for any of the EV
-                Chargers, or `None` otherwise.
         """
         evc_id = data.component_id
         new_state = EVChargerState.from_ev_charger_data(data)
-        if evc_id not in self._states or self._states[evc_id] != new_state:
-            self._states[evc_id] = new_state
-            return EVChargerPoolStates(self._states, evc_id)
-        return None
+        self._states[evc_id] = new_state
 
     async def _run(self) -> None:
         api_client = microgrid.connection_manager.get().api_client
         streams: list[Receiver[EVChargerData]] = await asyncio.gather(
             *[api_client.ev_charger_data(cid) for cid in self._component_ids]
         )
-
-        # Start with the `MISSING` state for all components.  This will change as data
-        # starts arriving from the individual components.
-        self._states = {
-            component_id: EVChargerState.MISSING for component_id in self._component_ids
-        }
         self._merged_stream = Merge(*streams)
-        sender = self._channel.new_sender()
-        await sender.send(self._get())
         async for data in self._merged_stream:
-            if updated_states := self._update(data):
-                await sender.send(updated_states)
-
-    def new_receiver(self) -> Receiver[EVChargerPoolStates]:
-        """Return a receiver that streams ev charger states.
-
-        Returns:
-            A receiver that streams the states of all EV Chargers in the pool, every
-                time the states of any of them change.
-        """
-        if self._task is None or self._task.done():
-            self._task = asyncio.create_task(self._run())
-        return self._channel.new_receiver()
+            self._update(data)
 
     async def stop(self) -> None:
         """Stop the status tracker."""
-        if self._task:
-            await cancel_and_await(self._task)
+        await cancel_and_await(self._task)
         if self._merged_stream:
             await self._merged_stream.stop()

tests/timeseries/test_ev_charger_pool.py

@@ -7,7 +7,6 @@
 
 import asyncio
 from math import isclose
-from typing import Optional
 
 from pytest_mock import MockerFixture
 
@@ -18,7 +17,6 @@
     EVChargerComponentState,
 )
 from frequenz.sdk.timeseries.ev_charger_pool._state_tracker import (
-    EVChargerPoolStates,
     EVChargerState,
     StateTracker,
 )
@@ -40,45 +38,39 @@ async def test_state_updates(self, mocker: MockerFixture) -> None:
         await mockgrid.start(mocker)
 
         state_tracker = StateTracker(set(mockgrid.evc_ids))
-        states = state_tracker.new_receiver()
 
-        async def check_next_state(
+        async def check_states(
             expected: dict[int, EVChargerState],
-            latest: Optional[tuple[int, EVChargerState]],
-        ) -> EVChargerPoolStates:
-            pool_states = await states.receive()
-            assert pool_states.latest_change() == latest
-            assert pool_states._states == expected  # pylint: disable=protected-access
-            return pool_states
+        ) -> None:
+            await asyncio.sleep(0.02)
+            for comp_id, exp_state in expected.items():
+                assert state_tracker.get(comp_id) == exp_state
 
         ## check that all chargers are in idle state.
         expected_states = {evc_id: EVChargerState.IDLE for evc_id in mockgrid.evc_ids}
         assert len(expected_states) == 5
-        await check_next_state(expected_states, None)
+        await check_states(expected_states)
 
         ## check that EV_PLUGGED state gets set
-        await asyncio.sleep(0.02)
         evc_2_id = mockgrid.evc_ids[2]
         mockgrid.evc_cable_states[evc_2_id] = EVChargerCableState.EV_PLUGGED
         mockgrid.evc_component_states[evc_2_id] = EVChargerComponentState.READY
         expected_states[evc_2_id] = EVChargerState.EV_PLUGGED
-        await check_next_state(expected_states, (evc_2_id, EVChargerState.EV_PLUGGED))
+        await check_states(expected_states)
 
         ## check that EV_LOCKED state gets set
-        await asyncio.sleep(0.03)
         evc_3_id = mockgrid.evc_ids[3]
         mockgrid.evc_cable_states[evc_3_id] = EVChargerCableState.EV_LOCKED
         mockgrid.evc_component_states[evc_3_id] = EVChargerComponentState.READY
         expected_states[evc_3_id] = EVChargerState.EV_LOCKED
-        await check_next_state(expected_states, (evc_3_id, EVChargerState.EV_LOCKED))
+        await check_states(expected_states)
 
         ## check that ERROR state gets set
-        await asyncio.sleep(0.1)
         evc_1_id = mockgrid.evc_ids[1]
         mockgrid.evc_cable_states[evc_1_id] = EVChargerCableState.EV_LOCKED
         mockgrid.evc_component_states[evc_1_id] = EVChargerComponentState.ERROR
         expected_states[evc_1_id] = EVChargerState.ERROR
-        await check_next_state(expected_states, (evc_1_id, EVChargerState.ERROR))
+        await check_states(expected_states)
 
         await state_tracker.stop()
         await mockgrid.cleanup()