diff --git a/deebot_client/capabilities.py b/deebot_client/capabilities.py index 5a794fac3..0b6a9ab12 100644 --- a/deebot_client/capabilities.py +++ b/deebot_client/capabilities.py @@ -125,6 +125,14 @@ class CapabilitySetTypes[E: Event, **P, T](CapabilitySet[E, P], CapabilityTypes[ """Capability for set command and types.""" +@dataclass(frozen=True, kw_only=True) +class CapabilityNumber[E: Event, **P](CapabilitySet[E, P]): + """Capability for a number entity with min and max.""" + + min: int + max: int + + @dataclass(frozen=True, kw_only=True) class CapabilityCleanAction: """Capabilities for clean action.""" @@ -228,11 +236,14 @@ class CapabilityStation: class CapabilityWater: """Capabilities for water.""" - amount: CapabilitySetTypes[ - water_info.WaterAmountEvent, - [water_info.WaterAmount | str], - water_info.WaterAmount, - ] + amount: ( + CapabilitySetTypes[ + water_info.WaterAmountEvent, + [water_info.WaterAmount | str], + water_info.WaterAmount, + ] + | CapabilityNumber[water_info.WaterCustomAmountEvent, [int]] + ) mop_attached: CapabilityEvent[water_info.MopAttachedEvent] diff --git a/deebot_client/commands/json/water_info.py b/deebot_client/commands/json/water_info.py index 107d05aa0..b74e119c3 100644 --- a/deebot_client/commands/json/water_info.py +++ b/deebot_client/commands/json/water_info.py @@ -11,6 +11,7 @@ SweepType, WaterAmount, WaterAmountEvent, + WaterCustomAmountEvent, WaterSweepTypeEvent, ) from deebot_client.message import HandlingResult @@ -35,7 +36,11 @@ def _handle_body_data_dict( :return: A message response """ - event_bus.notify(WaterAmountEvent(WaterAmount(int(data["amount"])))) + if "amount" in data: + event_bus.notify(WaterAmountEvent(WaterAmount(int(data["amount"])))) + + if "customAmount" in data: + event_bus.notify(WaterCustomAmountEvent(int(data["customAmount"]))) if (mop_attached := data.get("enable")) is not None: event_bus.notify(MopAttachedEvent(bool(mop_attached))) @@ -53,19 +58,32 @@ class SetWaterInfo(JsonSetCommand): get_command = GetWaterInfo _mqtt_params = MappingProxyType( { - "amount": InitParam(WaterAmount), + "amount": InitParam(WaterAmount, optional=True), + "customAmount": InitParam(int, "custom_amount", optional=True), "enable": None, # Remove it as we don't can set it (App includes it) "sweepType": InitParam(SweepType, "sweep_type", optional=True), } ) def __init__( - self, amount: WaterAmount | str, sweep_type: SweepType | str | None = None + self, + amount: WaterAmount | str | None = None, + custom_amount: int | None = None, + sweep_type: SweepType | str | None = None, ) -> None: params = {} - if isinstance(amount, str): - amount = get_enum(WaterAmount, amount) - params["amount"] = amount.value + if amount is not None: + if custom_amount is not None: + raise ValueError("Only one of amount or custom_amount can be provided.") + + if isinstance(amount, str): + amount = get_enum(WaterAmount, amount) + params["amount"] = amount.value + elif custom_amount is not None: + params["customAmount"] = custom_amount + else: + raise ValueError("Either amount or custom_amount must be provided.") + if sweep_type: if isinstance(sweep_type, str): sweep_type = get_enum(SweepType, sweep_type) diff --git a/deebot_client/events/water_info.py b/deebot_client/events/water_info.py index d0db80c03..6791cc9c2 100644 --- a/deebot_client/events/water_info.py +++ b/deebot_client/events/water_info.py @@ -11,6 +11,7 @@ "SweepType", "WaterAmount", "WaterAmountEvent", + "WaterCustomAmountEvent", "WaterSweepTypeEvent", ] @@ -37,6 +38,10 @@ class WaterAmountEvent(ValueEvent[WaterAmount]): """Water amount event.""" +class WaterCustomAmountEvent(ValueEvent[int]): + """Water custom amount event.""" + + class WaterSweepTypeEvent(ValueEvent[SweepType]): """Water sweep type event.""" diff --git a/deebot_client/hardware/deebot/xco2fc.py b/deebot_client/hardware/deebot/xco2fc.py index f9574d7bb..d168aec29 100644 --- a/deebot_client/hardware/deebot/xco2fc.py +++ b/deebot_client/hardware/deebot/xco2fc.py @@ -11,6 +11,7 @@ CapabilityExecute, CapabilityLifeSpan, CapabilityMap, + CapabilityNumber, CapabilitySet, CapabilitySetEnable, CapabilitySettings, @@ -190,16 +191,12 @@ total=CapabilityEvent(TotalStatsEvent, [GetTotalStats()]), ), water=CapabilityWater( - amount=CapabilitySetTypes( - event=water_info.WaterAmountEvent, + amount=CapabilityNumber( + event=water_info.WaterCustomAmountEvent, get=[GetWaterInfo()], - set=SetWaterInfo, - types=( - water_info.WaterAmount.LOW, - water_info.WaterAmount.MEDIUM, - water_info.WaterAmount.HIGH, - water_info.WaterAmount.ULTRAHIGH, - ), + set=lambda custom_amount: SetWaterInfo(custom_amount=custom_amount), + min=0, + max=50, ), mop_attached=CapabilityEvent(water_info.MopAttachedEvent, [GetWaterInfo()]), ), diff --git a/tests/commands/json/test_water_info.py b/tests/commands/json/test_water_info.py index 3b0826d7d..274a74ac0 100644 --- a/tests/commands/json/test_water_info.py +++ b/tests/commands/json/test_water_info.py @@ -11,6 +11,7 @@ SweepType, WaterAmount, WaterAmountEvent, + WaterCustomAmountEvent, WaterSweepTypeEvent, ) from tests.helpers import ( @@ -58,6 +59,21 @@ WaterSweepTypeEvent(SweepType.DEEP), ), ), + ( + { + "customAmount": 30, + "enable": 1, + "mopCount": 2, + "sideMop": 0, + "sweepType": 1, + "type": 1, + }, + ( + WaterCustomAmountEvent(30), + MopAttachedEvent(True), + WaterSweepTypeEvent(SweepType.STANDARD), + ), + ), ], ) async def test_GetWaterInfo(json: dict[str, Any], expected: tuple[Event, ...]) -> None: @@ -65,19 +81,52 @@ async def test_GetWaterInfo(json: dict[str, Any], expected: tuple[Event, ...]) - await assert_command(GetWaterInfo(), json, (firmware_event, *expected)) -@pytest.mark.parametrize(("water_value"), [WaterAmount.MEDIUM, "medium"]) -@pytest.mark.parametrize(("sweep_value"), [SweepType.STANDARD, "standard", None]) -async def test_SetWaterInfo_Wateramount( - water_value: WaterAmount | str, sweep_value: SweepType | str | None +@pytest.mark.parametrize( + ("command", "args", "expected_events"), + [ + ( + SetWaterInfo(WaterAmount.MEDIUM), + {"amount": 2}, + [WaterAmountEvent(WaterAmount.MEDIUM)], + ), + (SetWaterInfo("high"), {"amount": 3}, [WaterAmountEvent(WaterAmount.HIGH)]), + ( + SetWaterInfo(WaterAmount.LOW, sweep_type=SweepType.STANDARD), + {"amount": 1, "sweepType": 1}, + [ + WaterAmountEvent(WaterAmount.LOW), + WaterSweepTypeEvent(SweepType.STANDARD), + ], + ), + ( + SetWaterInfo(WaterAmount.ULTRAHIGH, sweep_type="deep"), + {"amount": 4, "sweepType": 2}, + [ + WaterAmountEvent(WaterAmount.ULTRAHIGH), + WaterSweepTypeEvent(SweepType.DEEP), + ], + ), + ( + SetWaterInfo(custom_amount=30), + {"customAmount": 30}, + [WaterCustomAmountEvent(30)], + ), + ( + SetWaterInfo(custom_amount=30, sweep_type="deep"), + {"customAmount": 30, "sweepType": 2}, + [ + WaterCustomAmountEvent(30), + WaterSweepTypeEvent(SweepType.DEEP), + ], + ), + ], +) +async def test_SetWaterInfo( + command: SetWaterInfo, + args: dict[str, Any], + expected_events: list[Event], ) -> None: - command = SetWaterInfo(water_value, sweep_value) - args = {"amount": 2} - expected_events: list[Event] = [ - WaterAmountEvent(WaterAmount.MEDIUM), - ] - if sweep_value: - args["sweepType"] = 1 - expected_events.append(WaterSweepTypeEvent(SweepType.STANDARD)) + """Test SetWaterInfo.""" await assert_set_command(command, args, expected_events) @@ -101,9 +150,19 @@ async def test_SetWaterInfo_Wateramount( ValueError, "'INEXSTING' is not a valid SweepType member", ), + ( + {}, + ValueError, + "Either amount or custom_amount must be provided.", + ), + ( + {"amount": WaterAmount.ULTRAHIGH, "custom_amount": "40"}, + ValueError, + "Only one of amount or custom_amount can be provided.", + ), ], ) -def test_SetWaterInfo_inexisting_value( +def test_SetWaterInfo_invalid( command_values: dict[str, Any], error: type[Exception], error_message: str ) -> None: with pytest.raises(error, match=error_message):