Skip to content

Commit 985a4a6

Browse files
shsmsshsms
committed
Extract microgrid mocking wrapper into a separate module
This can then be used in other tests that need a microgrid mock. Signed-off-by: Sahas Subramanian <[email protected]>
1 parent 065d631 commit 985a4a6

File tree

2 files changed

+105
-87
lines changed

2 files changed

+105
-87
lines changed

‎tests/actor/power_distributing/test_power_distributing.py‎

Lines changed: 3 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -6,21 +6,16 @@
66
from __future__ import annotations
77

88
import asyncio
9-
import dataclasses
109
import math
1110
import re
1211
from collections import abc
13-
from contextlib import asynccontextmanager
14-
from datetime import timedelta
15-
from typing import AsyncIterator, TypeVar
12+
from typing import TypeVar
1613
from unittest.mock import MagicMock
1714

18-
from frequenz.channels import Broadcast, Sender
15+
from frequenz.channels import Broadcast
1916
from frequenz.client.microgrid import ComponentCategory
2017
from pytest_mock import MockerFixture
2118

22-
from frequenz.sdk import microgrid
23-
from frequenz.sdk.actor import ResamplerConfig
2419
from frequenz.sdk.actor.power_distributing import (
2520
ComponentPoolStatus,
2621
PowerDistributingActor,
@@ -40,12 +35,11 @@
4035
Result,
4136
Success,
4237
)
43-
from frequenz.sdk.microgrid.component_graph import _MicrogridComponentGraph
4438
from frequenz.sdk.timeseries import Power
4539

4640
from ...conftest import SAFETY_TIMEOUT
41+
from ...microgrid.fixtures import _Mocks, _mocks
4742
from ...timeseries.mock_microgrid import MockMicrogrid
48-
from ...utils.component_data_streamer import MockComponentDataStreamer
4943
from ...utils.graph_generator import GraphGenerator
5044
from .test_battery_distribution_algorithm import (
5145
Bound,
@@ -60,84 +54,6 @@
6054
T = TypeVar("T") # Declare type variable
6155

6256

63-
@dataclasses.dataclass(frozen=True)
64-
class _Mocks:
65-
"""Mocks for the tests."""
66-
67-
microgrid: MockMicrogrid
68-
"""A mock microgrid instance."""
69-
70-
streamer: MockComponentDataStreamer
71-
"""A mock component data streamer."""
72-
73-
component_status_sender: Sender[ComponentPoolStatus]
74-
"""Sender for sending status of the components being tested."""
75-
76-
@classmethod
77-
async def new(
78-
cls,
79-
component_category: ComponentCategory,
80-
mocker: MockerFixture,
81-
graph: _MicrogridComponentGraph | None = None,
82-
grid_meter: bool | None = None,
83-
) -> _Mocks:
84-
"""Initialize the mocks."""
85-
mockgrid = MockMicrogrid(graph=graph, grid_meter=grid_meter, mocker=mocker)
86-
if not graph:
87-
mockgrid.add_batteries(3)
88-
await mockgrid.start()
89-
90-
# pylint: disable=protected-access
91-
if microgrid._data_pipeline._DATA_PIPELINE is not None:
92-
microgrid._data_pipeline._DATA_PIPELINE = None
93-
await microgrid._data_pipeline.initialize(
94-
ResamplerConfig(resampling_period=timedelta(seconds=0.1))
95-
)
96-
streamer = MockComponentDataStreamer(mockgrid.mock_client)
97-
98-
dp = microgrid._data_pipeline._DATA_PIPELINE
99-
assert dp is not None
100-
101-
if component_category == ComponentCategory.BATTERY:
102-
return cls(
103-
mockgrid,
104-
streamer,
105-
dp._battery_power_wrapper.status_channel.new_sender(),
106-
)
107-
raise ValueError(f"Unsupported component category: {component_category}")
108-
109-
async def stop(self) -> None:
110-
"""Stop the mocks."""
111-
# pylint: disable=protected-access
112-
assert microgrid._data_pipeline._DATA_PIPELINE is not None
113-
await asyncio.gather(
114-
*[
115-
microgrid._data_pipeline._DATA_PIPELINE._stop(),
116-
self.streamer.stop(),
117-
self.microgrid.cleanup(),
118-
]
119-
)
120-
# pylint: enable=protected-access
121-
122-
123-
@asynccontextmanager
124-
async def _mocks(
125-
mocker: MockerFixture,
126-
component_category: ComponentCategory,
127-
*,
128-
graph: _MicrogridComponentGraph | None = None,
129-
grid_meter: bool | None = None,
130-
) -> AsyncIterator[_Mocks]:
131-
"""Initialize the mocks."""
132-
mocks = await _Mocks.new(
133-
component_category, mocker, graph=graph, grid_meter=grid_meter
134-
)
135-
try:
136-
yield mocks
137-
finally:
138-
await mocks.stop()
139-
140-
14157
class TestPowerDistributingActor:
14258
# pylint: disable=protected-access
14359
# pylint: disable=too-many-public-methods

‎tests/microgrid/fixtures.py‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
# License: MIT
2+
# Copyright © 2023 Frequenz Energy-as-a-Service GmbH
3+
4+
"""Mock fixture for mocking the microgrid along with streaming."""
5+
6+
from __future__ import annotations
7+
8+
import asyncio
9+
import dataclasses
10+
from contextlib import asynccontextmanager
11+
from datetime import timedelta
12+
from typing import AsyncIterator
13+
14+
from frequenz.channels import Sender
15+
from frequenz.client.microgrid import ComponentCategory
16+
from pytest_mock import MockerFixture
17+
18+
from frequenz.sdk import microgrid
19+
from frequenz.sdk.actor import ResamplerConfig
20+
from frequenz.sdk.actor.power_distributing import ComponentPoolStatus
21+
from frequenz.sdk.microgrid.component_graph import _MicrogridComponentGraph
22+
23+
from ..timeseries.mock_microgrid import MockMicrogrid
24+
from ..utils.component_data_streamer import MockComponentDataStreamer
25+
26+
27+
@dataclasses.dataclass(frozen=True)
28+
class _Mocks:
29+
"""Mocks for the tests."""
30+
31+
microgrid: MockMicrogrid
32+
"""A mock microgrid instance."""
33+
34+
streamer: MockComponentDataStreamer
35+
"""A mock component data streamer."""
36+
37+
component_status_sender: Sender[ComponentPoolStatus]
38+
"""Sender for sending status of the components being tested."""
39+
40+
@classmethod
41+
async def new(
42+
cls,
43+
component_category: ComponentCategory,
44+
mocker: MockerFixture,
45+
graph: _MicrogridComponentGraph | None = None,
46+
grid_meter: bool | None = None,
47+
) -> _Mocks:
48+
"""Initialize the mocks."""
49+
mockgrid = MockMicrogrid(graph=graph, grid_meter=grid_meter, mocker=mocker)
50+
if not graph:
51+
mockgrid.add_batteries(3)
52+
await mockgrid.start()
53+
54+
# pylint: disable=protected-access
55+
if microgrid._data_pipeline._DATA_PIPELINE is not None:
56+
microgrid._data_pipeline._DATA_PIPELINE = None
57+
await microgrid._data_pipeline.initialize(
58+
ResamplerConfig(resampling_period=timedelta(seconds=0.1))
59+
)
60+
streamer = MockComponentDataStreamer(mockgrid.mock_client)
61+
62+
dp = microgrid._data_pipeline._DATA_PIPELINE
63+
assert dp is not None
64+
65+
if component_category == ComponentCategory.BATTERY:
66+
return cls(
67+
mockgrid,
68+
streamer,
69+
dp._battery_power_wrapper.status_channel.new_sender(),
70+
)
71+
raise ValueError(f"Unsupported component category: {component_category}")
72+
73+
async def stop(self) -> None:
74+
"""Stop the mocks."""
75+
# pylint: disable=protected-access
76+
assert microgrid._data_pipeline._DATA_PIPELINE is not None
77+
await asyncio.gather(
78+
*[
79+
microgrid._data_pipeline._DATA_PIPELINE._stop(),
80+
self.streamer.stop(),
81+
self.microgrid.cleanup(),
82+
]
83+
)
84+
# pylint: enable=protected-access
85+
86+
87+
@asynccontextmanager
88+
async def _mocks(
89+
mocker: MockerFixture,
90+
component_category: ComponentCategory,
91+
*,
92+
graph: _MicrogridComponentGraph | None = None,
93+
grid_meter: bool | None = None,
94+
) -> AsyncIterator[_Mocks]:
95+
"""Initialize the mocks."""
96+
mocks = await _Mocks.new(
97+
component_category, mocker, graph=graph, grid_meter=grid_meter
98+
)
99+
try:
100+
yield mocks
101+
finally:
102+
await mocks.stop()

0 commit comments

Comments
 (0)