Fix EZSP v14+ leaveNetwork and setManufacturerCode schemas

bellows/ezsp/__init__.py

@@ -316,12 +316,18 @@ def cb(frame_name, response):
         0,
     )
 
-    async def leaveNetwork(self, timeout: float | int = NETWORK_OPS_TIMEOUT) -> None:
+    async def leaveNetwork(
+        self,
+        timeout: float | int = NETWORK_OPS_TIMEOUT,
+        options: t.SlZigbeeLeaveNetworkOption = (
+            t.SlZigbeeLeaveNetworkOption.WITH_NO_OPTION
+        ),
+    ) -> None:
         """Send leaveNetwork command and wait for stackStatusHandler frame."""
         stack_status = asyncio.Future()
 
         with self.wait_for_stack_status(t.sl_Status.NETWORK_DOWN) as stack_status:
-            (status,) = await self._command("leaveNetwork")
+            status = await self._protocol.leave_network(options=options)
             if status != t.sl_Status.OK:
                 raise EzspError(f"failed to leave network: {status.name}")
 

bellows/ezsp/protocol.py

@@ -336,6 +336,15 @@ async def initialize_network(self) -> t.sl_Status:
     async def factory_reset(self) -> None:
         raise NotImplementedError
 
+    @abc.abstractmethod
+    async def leave_network(
+        self,
+        options: t.SlZigbeeLeaveNetworkOption = (
+            t.SlZigbeeLeaveNetworkOption.WITH_NO_OPTION
+        ),
+    ) -> t.sl_Status:
+        raise NotImplementedError
+
     @abc.abstractmethod
     async def send_unicast(
         self,

bellows/ezsp/v14/__init__.py

@@ -87,6 +87,15 @@ async def get_tc_link_key(self) -> zigpy.state.Key:
 
         return zigpy.state.Key(key=tc_link_key_data)
 
+    async def leave_network(
+        self,
+        options: t.SlZigbeeLeaveNetworkOption = (
+            t.SlZigbeeLeaveNetworkOption.WITH_NO_OPTION
+        ),
+    ) -> t.sl_Status:
+        (status,) = await self.leaveNetwork(options=options)
+        return status
+
     async def send_unicast(
         self,
         nwk: t.NWK,

bellows/ezsp/v14/commands.py

@@ -161,6 +161,26 @@ class GetTokenDataRsp(Struct):
             "status": t.EmberStatus,
         },
     ),
+    # Gained an `options` argument, which the NCP reads even if we don't send it
+    "leaveNetwork": (
+        0x0020,
+        {
+            "options": t.SlZigbeeLeaveNetworkOption,
+        },
+        {
+            "status": t.sl_Status,
+        },
+    ),
+    # Gained an `sl_status_t` response that `_REPLACEMENTS` can't add for us
+    "setManufacturerCode": (
+        0x0015,
+        {
+            "code": t.uint16_t,
+        },
+        {
+            "status": t.sl_Status,
+        },
+    ),
 }
 
 

bellows/ezsp/v4/__init__.py

@@ -127,6 +127,18 @@ async def initialize_network(self) -> t.sl_Status:
     async def factory_reset(self) -> None:
         await self.clearKeyTable()
 
+    async def leave_network(
+        self,
+        options: t.SlZigbeeLeaveNetworkOption = (
+            t.SlZigbeeLeaveNetworkOption.WITH_NO_OPTION
+        ),
+    ) -> t.sl_Status:
+        if options != t.SlZigbeeLeaveNetworkOption.WITH_NO_OPTION:
+            raise ValueError(f"Leave options are not supported: {options!r}")
+
+        (status,) = await self.leaveNetwork()
+        return t.sl_Status.from_ember_status(status)
+
     async def send_unicast(
         self,
         nwk: t.NWK,

bellows/types/named.py

@@ -2889,6 +2889,14 @@ class RouteRecordConcentratortype(basic.enum8):
     HIGH_RAM = 2
 
 
+class SlZigbeeLeaveNetworkOption(basic.bitmap8):
+    """Options for leaving the network (EZSP v14+)."""
+
+    WITH_NO_OPTION = 0x00
+    WITH_OPTION_REJOIN = 0x20
+    IS_REQUESTED = 0x40
+
+
 class SlZigbeeDhcPaAlgorithm(basic.enum8):
     """Dynamic Hardware Configuration PA descriptor algorithm."""
 

tests/test_ezsp.py

@@ -451,33 +451,47 @@ async def test_set_disable_source_routing(ezsp_f):
 async def test_leave_network_error(ezsp_f):
     """Test EZSP leaveNetwork command failure."""
 
-    with patch.object(ezsp_f, "_command", new_callable=AsyncMock) as cmd_mock:
-        cmd_mock.return_value = [t.EmberStatus.ERR_FATAL]
+    with patch.object(ezsp_f._protocol, "leave_network", new_callable=AsyncMock) as m:
+        m.return_value = t.sl_Status.FAIL
         with pytest.raises(EzspError):
             await ezsp_f.leaveNetwork(timeout=0.01)
 
 
 async def test_leave_network_no_stack_status(ezsp_f):
     """Test EZSP leaveNetwork command, no stackStatusHandler callback."""
 
-    with patch.object(ezsp_f, "_command", new_callable=AsyncMock) as cmd_mock:
-        cmd_mock.return_value = [t.EmberStatus.SUCCESS]
+    with patch.object(ezsp_f._protocol, "leave_network", new_callable=AsyncMock) as m:
+        m.return_value = t.sl_Status.OK
         with pytest.raises(TimeoutError):
             await ezsp_f.leaveNetwork(timeout=0.01)
 
 
-async def test_leave_network(ezsp_f):
+@pytest.mark.parametrize(
+    ("kwargs", "options"),
+    [
+        ({}, t.SlZigbeeLeaveNetworkOption.WITH_NO_OPTION),
+        (
+            {"options": t.SlZigbeeLeaveNetworkOption.WITH_OPTION_REJOIN},
+            t.SlZigbeeLeaveNetworkOption.WITH_OPTION_REJOIN,
+        ),
+    ],
+)
+async def test_leave_network(
+    ezsp_f, kwargs: dict, options: t.SlZigbeeLeaveNetworkOption
+) -> None:
     """Test EZSP leaveNetwork command."""
 
-    async def _mock_cmd(*args, **kwargs):
+    async def _mock_leave(*args, **kwargs):
         ezsp_f.handle_callback("stackStatusHandler", [t.EmberStatus.NETWORK_UP])
         ezsp_f.handle_callback("stackStatusHandler", [t.EmberStatus.NETWORK_UP])
         ezsp_f.handle_callback("stackStatusHandler", [t.EmberStatus.NETWORK_DOWN])
-        return [t.EmberStatus.SUCCESS]
+        return t.sl_Status.OK
+
+    with patch.object(ezsp_f._protocol, "leave_network", new_callable=AsyncMock) as m:
+        m.side_effect = _mock_leave
+        await ezsp_f.leaveNetwork(timeout=0.01, **kwargs)
 
-    with patch.object(ezsp_f, "_command", new_callable=AsyncMock) as cmd_mock:
-        cmd_mock.side_effect = _mock_cmd
-        await ezsp_f.leaveNetwork(timeout=0.01)
+    assert m.mock_calls == [call(options=options)]
 
 
 async def test_xncp_token_override(ezsp_f):

tests/test_ezsp_v14.py

@@ -4,6 +4,7 @@
 import zigpy.exceptions
 import zigpy.state
 
+import bellows.ezsp
 import bellows.ezsp.v14
 import bellows.types as t
 
@@ -33,6 +34,64 @@ def test_ezsp_frame_rx(ezsp_f):
     assert ezsp_f._handle_callback.call_args[0][1] == [0x01, 0x02, 0x1234]
 
 
+@pytest.mark.parametrize(
+    "version",
+    [v for v in bellows.ezsp.EZSP._BY_VERSION if v >= 14],
+)
+def test_leave_network_tx_schema(version: int) -> None:
+    """`leaveNetwork` sends its one-byte `options` argument on v14 and newer."""
+    ezsp = bellows.ezsp.EZSP._BY_VERSION[version](MagicMock(), MagicMock())
+    ezsp._seq = 0x22
+
+    data = ezsp._ezsp_frame(
+        "leaveNetwork", options=t.SlZigbeeLeaveNetworkOption.WITH_NO_OPTION
+    )
+    assert data == b"\x22\x00\x01\x20\x00" + b"\x00"
+
+
+@pytest.mark.parametrize(
+    "version",
+    [v for v in bellows.ezsp.EZSP._BY_VERSION if v >= 14],
+)
+@pytest.mark.parametrize(
+    ("command", "data", "expected"),
+    [
+        ("leaveNetwork", b"\x00\x00\x00\x00", [t.sl_Status.OK]),
+        ("setManufacturerCode", b"\x00\x00\x00\x00", [t.sl_Status.OK]),
+    ],
+)
+def test_status_rx_schemas(
+    version: int, command: str, data: bytes, expected: list
+) -> None:
+    """The `sl_status_t` responses are fully parsed by v14 and newer."""
+    _, _, rx_schema = bellows.ezsp.EZSP._BY_VERSION[version].COMMANDS[command]
+    result, rest = t.deserialize_dict(data, rx_schema)
+
+    assert list(result.values()) == expected
+    assert rest == b""
+
+
+async def test_leave_network(ezsp_f) -> None:
+    ezsp_f.leaveNetwork.return_value = (t.sl_Status.OK,)
+    assert await ezsp_f.leave_network() == t.sl_Status.OK
+    assert ezsp_f.leaveNetwork.mock_calls == [
+        call(options=t.SlZigbeeLeaveNetworkOption.WITH_NO_OPTION)
+    ]
+
+
+async def test_leave_network_options(ezsp_f) -> None:
+    ezsp_f.leaveNetwork.return_value = (t.sl_Status.OK,)
+    assert (
+        await ezsp_f.leave_network(
+            options=t.SlZigbeeLeaveNetworkOption.WITH_OPTION_REJOIN
+        )
+        == t.sl_Status.OK
+    )
+    assert ezsp_f.leaveNetwork.mock_calls == [
+        call(options=t.SlZigbeeLeaveNetworkOption.WITH_OPTION_REJOIN)
+    ]
+
+
 async def test_read_address_table(ezsp_f):
     def get_addr_table_info(index):
         default = (

tests/test_ezsp_v4.py

@@ -277,6 +277,28 @@ async def test_factory_reset(ezsp_f) -> None:
     assert ezsp_f.clearKeyTable.mock_calls == [call()]
 
 
+async def test_leave_network(ezsp_f) -> None:
+    ezsp_f.leaveNetwork.return_value = (t.EmberStatus.SUCCESS,)
+    assert await ezsp_f.leave_network() == t.sl_Status.OK
+    assert ezsp_f.leaveNetwork.mock_calls == [call()]
+
+
+@pytest.mark.parametrize(
+    "options",
+    [
+        t.SlZigbeeLeaveNetworkOption.WITH_OPTION_REJOIN,
+        t.SlZigbeeLeaveNetworkOption.IS_REQUESTED,
+    ],
+)
+async def test_leave_network_options_unsupported(
+    ezsp_f, options: t.SlZigbeeLeaveNetworkOption
+) -> None:
+    with pytest.raises(ValueError):
+        await ezsp_f.leave_network(options=options)
+
+    assert ezsp_f.leaveNetwork.mock_calls == []
+
+
 async def test_send_unicast(ezsp_f) -> None:
     ezsp_f.sendUnicast.return_value = (t.EmberStatus.SUCCESS, 0x42)
     status, message_tag = await ezsp_f.send_unicast(