Add a test case for a composite formula with constants

This is required to simulate composing on formula builders. Earlier,
we were testing composing only on pre-built formula engines.

Also added some type hints to ensure they pass mypy checks.

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

Author: shsms

tests/timeseries/_formula_engine/test_formula_composition.py

@@ -12,7 +12,7 @@
 from pytest_mock import MockerFixture
 
 from frequenz.sdk import microgrid
-from frequenz.sdk.timeseries._quantities import Power
+from frequenz.sdk.timeseries import Power, Sample
 
 from ..mock_microgrid import MockMicrogrid
 from .utils import get_resampled_stream
@@ -303,7 +303,10 @@ async def test_formula_composition_min_max_const(
             max_pow = await engine_max_rx.receive()
             assert max_pow and max_pow.value and max_pow.value.isclose(Power.zero())
 
-    async def test_formula_composition_constant(self, mocker: MockerFixture) -> None:
+    async def test_formula_composition_constant(  # pylint: disable=too-many-locals
+        self,
+        mocker: MockerFixture,
+    ) -> None:
         """Test the composition of formulas with constant values."""
         async with (
             MockMicrogrid(grid_meter=True, mocker=mocker) as mockgrid,
@@ -334,40 +337,64 @@ async def test_formula_composition_constant(self, mocker: MockerFixture) -> None
                 engine_div._stop  # pylint: disable=protected-access
             )
 
+            engine_composite = (
+                (
+                    (grid.power + Power.from_watts(50.0)) / 2.0
+                    + grid.power
+                    - Power.from_watts(20.0)
+                )
+                * 2.0
+            ).build("grid_power_composite")
+
             await mockgrid.mock_resampler.send_meter_power([100.0])
 
             # Test addition
-            grid_power_addition = await engine_add.new_receiver().receive()
+            grid_power_addition: Sample[Power] = (
+                await engine_add.new_receiver().receive()
+            )
             assert grid_power_addition.value is not None
             assert math.isclose(
                 grid_power_addition.value.as_watts(),
                 150.0,
             )
 
             # Test subtraction
-            grid_power_subtraction = await engine_sub.new_receiver().receive()
+            grid_power_subtraction: Sample[Power] = (
+                await engine_sub.new_receiver().receive()
+            )
             assert grid_power_subtraction.value is not None
             assert math.isclose(
                 grid_power_subtraction.value.as_watts(),
                 0.0,
             )
 
             # Test multiplication
-            grid_power_multiplication = await engine_mul.new_receiver().receive()
+            grid_power_multiplication: Sample[Power] = (
+                await engine_mul.new_receiver().receive()
+            )
             assert grid_power_multiplication.value is not None
             assert math.isclose(
                 grid_power_multiplication.value.as_watts(),
                 200.0,
             )
 
             # Test division
-            grid_power_division = await engine_div.new_receiver().receive()
+            grid_power_division: Sample[Power] = (
+                await engine_div.new_receiver().receive()
+            )
             assert grid_power_division.value is not None
             assert math.isclose(
                 grid_power_division.value.as_watts(),
                 50.0,
             )
 
+            # Test composite formula
+            grid_power_composite: Sample[Power] = (
+                await engine_composite.new_receiver().receive()
+            )
+            assert grid_power_composite.value is not None
+            assert math.isclose(grid_power_composite.value.as_watts(), 310.0)
+
             # Test multiplication with a Quantity
             with pytest.raises(RuntimeError):
                 engine_assert = (grid.power * Power.from_watts(2.0)).build(  # type: ignore