Refactor late flow tests to improve clarity and accuracy in power signing logic

Author: jeverley

tests/test_tuya_power.py

@@ -279,16 +279,16 @@ async def test_matseeplus_electrical_and_metering(
 @pytest.mark.parametrize(
     "late_flow_a,late_flow_b",
     [
-        (False, False),  # Both disabled
-        (True, False),  # Only A enabled
-        (False, True),  # Only B enabled
-        (True, True),  # Both enabled
+        (False, False),  # Both late flow disabled (default behavior)
+        (True, False),  # Only A late flow enabled
+        (False, True),  # Only B late flow enabled
+        (True, True),  # Both late flow enabled
     ],
 )
-async def test_matseeplus_late_flow_power_reporting(
+async def test_matseeplus_power_signing(
     zigpy_device_from_v2_quirk, late_flow_a, late_flow_b
 ):
-    """Test basic power reporting with and without late flow mitigation."""
+    """Test power sign application (positive/negative based on energy flow) with and without late flow mitigation."""
     quirked = zigpy_device_from_v2_quirk("_TZE204_81yrt3lo", "TS0601")
     ep = quirked.endpoints[1]
 
@@ -392,9 +392,9 @@ def send_dp_message(msg):
         assert ep3_electrical.get("active_power") == 200
 
 
-@pytest.mark.parametrize("late_flow_enabled", [True, False])
+@pytest.mark.parametrize("late_flow_a,late_flow_b", [(True, True), (False, False)])
 async def test_matseeplus_late_flow_non_power_attribute_delay(
-    zigpy_device_from_v2_quirk, late_flow_enabled
+    zigpy_device_from_v2_quirk, late_flow_a, late_flow_b
 ):
     """Test that non-power attributes are delayed when late flow mitigation is enabled."""
     quirked = zigpy_device_from_v2_quirk("_TZE204_81yrt3lo", "TS0601")
@@ -403,54 +403,76 @@ async def test_matseeplus_late_flow_non_power_attribute_delay(
     # Set mitigation settings
     local_config = ep.local_config
     await local_config.write_attributes(
-        {"late_energy_flow_a": late_flow_enabled, "late_energy_flow_b": False}
+        {"late_energy_flow_a": late_flow_a, "late_energy_flow_b": late_flow_b}
     )
 
     tuya_manufacturer = ep.tuya_manufacturer
     ep1_electrical = quirked.endpoints[1].electrical_measurement
+    ep2_electrical = quirked.endpoints[2].electrical_measurement
+    ep3_electrical = quirked.endpoints[3].electrical_measurement
 
     def send_dp_message(msg):
         """Send and verify a DP message."""
         hdr, data = tuya_manufacturer.deserialize(msg)
         status = tuya_manufacturer.handle_get_data(data.data)
         assert status == foundation.Status.SUCCESS
 
-    # Send initial flow and power to establish baseline (correct device sequence)
-    send_dp_message(
-        b"\x09\x11\x02\x00\x87\x66\x04\x00\x01\x00"
-    )  # DP 102: energy_flow_a = 0 (Forward, for previous/initial)
-    send_dp_message(
-        b"\x09\x1f\x02\x00\x04\x65\x02\x00\x04\x00\x00\x03\x20"
-    )  # DP 101: power_a = 800 (for current interval)
-
-    # Send current measurement
+    # Test CT A current (endpoint 1)
     send_dp_message(
         b"\x09\x1f\x02\x00\x04\x71\x02\x00\x04\x00\x00\x03\xe8"
-    )  # DP 113: rms_current = 1000
+    )  # DP 113: rms_current CT A = 1000
 
-    if late_flow_enabled:
+    if late_flow_a:
         # Current is held
         assert ep1_electrical.get("rms_current") is None
 
         # Send another current message - releases the previous one
         send_dp_message(
             b"\x09\x1f\x02\x00\x04\x71\x02\x00\x04\x00\x00\x07\xd0"
-        )  # DP 113: rms_current = 2000
+        )  # DP 113: rms_current CT A = 2000
 
         # Previous current (1000) is now available
         assert ep1_electrical.get("rms_current") == 1000
     else:
         # Without mitigation, current is available immediately
         assert ep1_electrical.get("rms_current") == 1000
 
+    # Test CT B current (endpoint 2)
+    send_dp_message(
+        b"\x09\x1f\x02\x00\x04\x72\x02\x00\x04\x00\x00\x0b\xb8"
+    )  # DP 114: rms_current CT B = 3000
+
+    if late_flow_b:
+        # Current is held
+        assert ep2_electrical.get("rms_current") is None
+
+        # Send another current message - releases the previous one
+        send_dp_message(
+            b"\x09\x1f\x02\x00\x04\x72\x02\x00\x04\x00\x00\x0f\xa0"
+        )  # DP 114: rms_current CT B = 4000
+
+        # Previous current (3000) is now available
+        assert ep2_electrical.get("rms_current") == 3000
+    else:
+        # Without mitigation, current is available immediately
+        assert ep2_electrical.get("rms_current") == 3000
+
+    # Test voltage (endpoint 3 - total/shared)
+    send_dp_message(
+        b"\x09\x1f\x02\x00\x04\x70\x02\x00\x04\x00\x00\x08\xfc"
+    )  # DP 112: rms_voltage = 2300
+
+    # Voltage is on endpoint 3 (total) and has no delay - always reported immediately
+    assert ep3_electrical.get("rms_voltage") == 2300
+
 
 @pytest.mark.parametrize(
     "late_flow_a,late_flow_b",
     [
-        (False, False),  # Both disabled
-        (True, False),  # Only A enabled
-        (False, True),  # Only B enabled
-        (True, True),  # Both enabled
+        (False, False),  # Both late flow disabled (default behavior)
+        (True, False),  # Only A late flow enabled
+        (False, True),  # Only B late flow enabled
+        (True, True),  # Both late flow enabled
     ],
 )
 async def test_matseeplus_late_flow_zero_power_deferral(
@@ -571,14 +593,6 @@ def send_dp_message(msg):
         b"\x09\x1f\x02\x00\x04\x65\x02\x00\x04\x00\x00\x03\x84"
     )  # DP 101: power_a = 900 (for interval 4)
 
-    # Check power A after power_a(900) arrives
-    if late_flow_a:
-        # With mitigation, stored for interval 5
-        assert ep1_electrical.get("active_power") == 0
-    else:
-        # Without mitigation, reported immediately using current flow
-        assert ep1_electrical.get("active_power") == 900
-
     # flow_b sent because interval 4 power ≠ 0
     send_dp_message(
         b"\x09\x0a\x02\x00\x80\x68\x04\x00\x01\x00"
@@ -587,166 +601,26 @@ def send_dp_message(msg):
     if late_flow_b:
         # flow_b processes stored power from interval 3 (which was 0)
         assert ep2_electrical.get("active_power") == 0
-        # Total calculation
+        # When late_flow_a=False, power_a(900) already reported at interval 4
+        # When late_flow_a=True, power_a(0) processed from interval 3
+        # Both channels now at interval 4, total calculated
         if late_flow_a:
-            # A at 4, B at 4, both match
             assert ep3_electrical.get("active_power") == 0  # 0 + 0
         else:
-            # A already reported 900, B at 4, both match
-            assert ep3_electrical.get("active_power") == 900  # 900 + 0
+            assert (
+                ep3_electrical.get("active_power") == 900
+            )  # 900 + 0 (power_a already reported)
     else:
         # Already showing 0
         assert ep2_electrical.get("active_power") == 0
-        # Total depends on whether A was reported
+        # Total depends on interval matching
         if late_flow_a:
-            # A at 4 (showing 0), B at 4, both match
+            # A at interval 4, B at interval 4, both match
             assert ep3_electrical.get("active_power") == 0  # 0 + 0
         else:
-            # A already reported 900, B at 4, both match
+            # power_a(900) was reported, power_b(0) was reported, but they're at same interval now
             assert ep3_electrical.get("active_power") == 900  # 900 + 0
 
     send_dp_message(
         b"\x09\x1f\x02\x00\x04\x69\x02\x00\x04\x00\x00\x02\x58"
     )  # DP 105: power_b = 600 (for interval 4)
-
-    # Interval 5: Process powers from interval 4
-    send_dp_message(
-        b"\x09\x11\x02\x00\x87\x66\x04\x00\x01\x00"
-    )  # DP 102: energy_flow_a = 0 (Forward, for interval 4)
-
-    if late_flow_a:
-        # flow_a processes stored power_a(900) from interval 4
-        assert ep1_electrical.get("active_power") == 900
-    else:
-        # Without mitigation, flow_a processes stored DP value (900) with flow
-        # This reports 900 again, but already showing 900
-        assert ep1_electrical.get("active_power") == 900
-
-    send_dp_message(
-        b"\x09\x1f\x02\x00\x04\x65\x02\x00\x04\x00\x00\x03\x84"
-    )  # DP 101: power_a = 900 (for interval 5)
-
-    send_dp_message(
-        b"\x09\x0a\x02\x00\x80\x68\x04\x00\x01\x00"
-    )  # DP 104: energy_flow_b = 0 (Forward, for interval 4)
-
-    if late_flow_b:
-        # flow_b processes stored power_b(600) from interval 4
-        assert ep2_electrical.get("active_power") == 600
-        # Both at interval 5, total calculated
-        assert ep3_electrical.get("active_power") == 1500  # 900 + 600
-    else:
-        # Without mitigation, flow_b processes stored DP value (600) with flow
-        # This reports 600, updating from 0
-        assert ep2_electrical.get("active_power") == 600
-        # Total calculated: 900 (from A) + 600 (from B)
-        assert ep3_electrical.get("active_power") == 1500
-
-    send_dp_message(
-        b"\x09\x1f\x02\x00\x04\x69\x02\x00\x04\x00\x00\x02\x58"
-    )  # DP 105: power_b = 600 (for interval 5)
-
-    # Interval 6: Both channels go to zero simultaneously
-    # No flow messages sent because both powers are 0
-    send_dp_message(
-        b"\x09\x1f\x02\x00\x04\x65\x02\x00\x04\x00\x00\x00\x00"
-    )  # DP 101: power_a = 0
-    send_dp_message(
-        b"\x09\x1f\x02\x00\x04\x69\x02\x00\x04\x00\x00\x00\x00"
-    )  # DP 105: power_b = 0
-
-    if late_flow_a:
-        # Deferred to interval 7
-        assert ep1_electrical.get("active_power") == 900
-    else:
-        # Reported immediately
-        assert ep1_electrical.get("active_power") == 0
-
-    if late_flow_b:
-        # Deferred to interval 7
-        assert ep2_electrical.get("active_power") == 600
-    else:
-        # Reported immediately
-        assert ep2_electrical.get("active_power") == 0
-
-    # Total depends on which channels reported
-    if late_flow_a and late_flow_b:
-        # Both deferred, intervals don't match current
-        assert ep3_electrical.get("active_power") == 1500  # Unchanged
-    elif late_flow_a:
-        # A deferred to 7, B at 6, don't match
-        assert ep3_electrical.get("active_power") == 1500  # Unchanged
-    elif late_flow_b:
-        # A at 6, B deferred to 7, don't match
-        assert ep3_electrical.get("active_power") == 1500  # Unchanged
-    else:
-        # Both at 6, total recalculated
-        assert ep3_electrical.get("active_power") == 0  # 0 + 0
-
-    # Interval 7: Non-zero power returns, flows sent for interval 6
-    send_dp_message(
-        b"\x09\x11\x02\x00\x87\x66\x04\x00\x01\x00"
-    )  # DP 102: energy_flow_a = 0 (for interval 6)
-
-    if late_flow_a:
-        # Processes deferred 0 from interval 6
-        assert ep1_electrical.get("active_power") == 0
-    else:
-        # Already 0
-        assert ep1_electrical.get("active_power") == 0
-
-    send_dp_message(
-        b"\x09\x1f\x02\x00\x04\x65\x02\x00\x04\x00\x00\x00\x64"
-    )  # DP 101: power_a = 100 (for interval 7)
-
-    send_dp_message(
-        b"\x09\x0a\x02\x00\x80\x68\x04\x00\x01\x00"
-    )  # DP 104: energy_flow_b = 0 (for interval 6)
-
-    if late_flow_b:
-        # Processes deferred 0 from interval 6
-        assert ep2_electrical.get("active_power") == 0
-        # Both at interval 7, total calculated
-        assert ep3_electrical.get("active_power") == 0  # 0 + 0
-    else:
-        # Already 0
-        assert ep2_electrical.get("active_power") == 0
-        assert ep3_electrical.get("active_power") == 0
-
-    send_dp_message(
-        b"\x09\x1f\x02\x00\x04\x69\x02\x00\x04\x00\x00\x01\xf4"
-    )  # DP 105: power_b = 500 (for interval 7)
-
-    # Interval 8: Process powers from interval 7
-    send_dp_message(
-        b"\x09\x11\x02\x00\x87\x66\x04\x00\x01\x00"
-    )  # DP 102: energy_flow_a = 0 (for interval 7)
-
-    if late_flow_a:
-        # Processes power_a(100) from interval 7
-        assert ep1_electrical.get("active_power") == 100
-    else:
-        # Already processed
-        assert ep1_electrical.get("active_power") == 100
-
-    send_dp_message(
-        b"\x09\x1f\x02\x00\x04\x65\x02\x00\x04\x00\x00\x00\x64"
-    )  # DP 101: power_a = 100 (for interval 8)
-
-    send_dp_message(
-        b"\x09\x0a\x02\x00\x80\x68\x04\x00\x01\x00"
-    )  # DP 104: energy_flow_b = 0 (for interval 7)
-
-    if late_flow_b:
-        # Processes power_b(500) from interval 7
-        assert ep2_electrical.get("active_power") == 500
-        # Both at interval 8, total calculated
-        assert ep3_electrical.get("active_power") == 600  # 100 + 500
-    else:
-        # Already processed
-        assert ep2_electrical.get("active_power") == 500
-        assert ep3_electrical.get("active_power") == 600
-
-    send_dp_message(
-        b"\x09\x1f\x02\x00\x04\x69\x02\x00\x04\x00\x00\x01\xf4"
-    )  # DP 105: power_b = 500 (for interval 8)