@@ -114,6 +114,14 @@ async def _send_time_changed(zha_gateway: Gateway, seconds: int):
114114    "zha.zigbee.cluster_handlers.general.BasicClusterHandler.async_initialize" , 
115115    new = mock .AsyncMock (), 
116116) 
117+ @pytest .mark .parametrize ( 
118+     "zha_gateway" , 
119+     [ 
120+         "zha_gateway" , 
121+         "ws_gateways" , 
122+     ], 
123+     indirect = True , 
124+ ) 
117125async  def  test_check_available_success (
118126    zha_gateway : Gateway ,
119127    caplog : pytest .LogCaptureFixture ,
@@ -124,19 +132,28 @@ async def test_check_available_success(
124132    )
125133    zha_device  =  await  join_zigpy_device (zha_gateway , device_with_basic_cluster_handler )
126134    basic_ch  =  device_with_basic_cluster_handler .endpoints [3 ].basic 
135+     if  hasattr (zha_gateway , "ws_gateway" ):
136+         server_device  =  zha_gateway .ws_gateway .devices [zha_device .ieee ]
137+         server_gateway  =  zha_gateway .ws_gateway 
138+     else :
139+         server_device  =  zha_device 
140+         server_gateway  =  zha_gateway 
127141
128142    assert  not  zha_device .is_coordinator 
129143    assert  not  zha_device .is_active_coordinator 
130144
131145    basic_ch .read_attributes .reset_mock ()
132146    device_with_basic_cluster_handler .last_seen  =  None 
133147    assert  zha_device .available  is  True 
134-     await  _send_time_changed (zha_gateway , zha_device .consider_unavailable_time  +  2 )
148+     await  _send_time_changed (zha_gateway , server_device .consider_unavailable_time  +  2 )
135149    assert  zha_device .available  is  False 
136150    assert  basic_ch .read_attributes .await_count  ==  0 
137151
152+     for  entity  in  server_device .platform_entities .values ():
153+         assert  not  entity .available 
154+ 
138155    device_with_basic_cluster_handler .last_seen  =  (
139-         time .time () -  zha_device .consider_unavailable_time  -  100 
156+         time .time () -  server_device .consider_unavailable_time  -  100 
140157    )
141158    _seens  =  [time .time (), device_with_basic_cluster_handler .last_seen ]
142159
@@ -146,63 +163,82 @@ def _update_last_seen(*args, **kwargs):  # pylint: disable=unused-argument
146163
147164    basic_ch .read_attributes .side_effect  =  _update_last_seen 
148165
149-     for  entity  in  zha_device .platform_entities .values ():
166+     for  entity  in  server_device .platform_entities .values ():
150167        entity .emit  =  mock .MagicMock (wraps = entity .emit )
151168
152169    # we want to test the device availability handling alone 
153-     zha_gateway .global_updater .stop ()
170+     server_gateway .global_updater .stop ()
154171
155172    # successfully ping zigpy device, but zha_device is not yet available 
156173    await  _send_time_changed (
157-         zha_gateway , zha_gateway ._device_availability_checker .__polling_interval  +  1 
174+         zha_gateway , server_gateway ._device_availability_checker .__polling_interval  +  1 
158175    )
159176    assert  basic_ch .read_attributes .await_count  ==  1 
160177    assert  basic_ch .read_attributes .await_args [0 ][0 ] ==  ["manufacturer" ]
161178    assert  zha_device .available  is  False 
162179
163-     for  entity  in  zha_device .platform_entities .values ():
180+     for  entity  in  server_device .platform_entities .values ():
164181        entity .emit .assert_not_called ()
165182        assert  not  entity .available 
183+         if  server_device  !=  zha_device :
184+             assert  not  zha_device .platform_entities [
185+                 (entity .PLATFORM , entity .unique_id )
186+             ].available 
166187        entity .emit .reset_mock ()
167188
168189    # There was traffic from the device: pings, but not yet available 
169190    await  _send_time_changed (
170-         zha_gateway , zha_gateway ._device_availability_checker .__polling_interval  +  1 
191+         zha_gateway , server_gateway ._device_availability_checker .__polling_interval  +  1 
171192    )
172193    assert  basic_ch .read_attributes .await_count  ==  2 
173194    assert  basic_ch .read_attributes .await_args [0 ][0 ] ==  ["manufacturer" ]
174195    assert  zha_device .available  is  False 
175196
176-     for  entity  in  zha_device .platform_entities .values ():
197+     for  entity  in  server_device .platform_entities .values ():
177198        entity .emit .assert_not_called ()
178199        assert  not  entity .available 
200+         if  server_device  !=  zha_device :
201+             assert  not  zha_device .platform_entities [
202+                 (entity .PLATFORM , entity .unique_id )
203+             ].available 
179204        entity .emit .reset_mock ()
180205
181206    # There was traffic from the device: don't try to ping, marked as available 
182207    await  _send_time_changed (
183-         zha_gateway , zha_gateway ._device_availability_checker .__polling_interval  +  1 
208+         zha_gateway , server_gateway ._device_availability_checker .__polling_interval  +  1 
184209    )
185210    assert  basic_ch .read_attributes .await_count  ==  2 
186211    assert  basic_ch .read_attributes .await_args [0 ][0 ] ==  ["manufacturer" ]
187212    assert  zha_device .available  is  True 
188213    assert  zha_device .on_network  is  True 
189214
190-     for  entity  in  zha_device .platform_entities .values ():
215+     for  entity  in  server_device .platform_entities .values ():
191216        entity .emit .assert_called ()
217+         if  server_device  !=  zha_device :
218+             assert  zha_device .platform_entities [
219+                 (entity .PLATFORM , entity .unique_id )
220+             ].available 
192221        assert  entity .available 
193222        entity .emit .reset_mock ()
194223
195224    assert  "Device is not on the network, marking unavailable"  not  in caplog .text 
196-     zha_device .on_network  =  False 
225+     server_gateway ._device_availability_checker .stop ()
226+ 
227+     server_device .on_network  =  False 
228+     await  zha_gateway .async_block_till_done (wait_background_tasks = True )
197229
198230    assert  zha_device .available  is  False 
199231    assert  zha_device .on_network  is  False 
200232
201233    assert  "Device is not on the network, marking unavailable"  in  caplog .text 
202234
203-     for  entity  in  zha_device .platform_entities .values ():
235+     for  entity  in  server_device .platform_entities .values ():
204236        entity .emit .assert_called ()
205237        assert  not  entity .available 
238+         if  server_device  !=  zha_device :
239+             assert  not  zha_device .platform_entities [
240+                 (entity .PLATFORM , entity .unique_id )
241+             ].available 
206242        entity .emit .reset_mock ()
207243
208244
0 commit comments