async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: """Set up the Ambient PWS as config entry.""" if not entry.unique_id: hass.config_entries.async_update_entry( entry, unique_id=entry.data[CONF_APP_KEY]) ambient = AmbientStation( hass, entry, Websocket(entry.data[CONF_APP_KEY], entry.data[CONF_API_KEY]), ) try: await ambient.ws_connect() except WebsocketError as err: LOGGER.error("Config entry failed: %s", err) raise ConfigEntryNotReady from err hass.data.setdefault(DOMAIN, {})[entry.entry_id] = ambient async def _async_disconnect_websocket(_: Event) -> None: await ambient.websocket.disconnect() entry.async_on_unload( hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_disconnect_websocket)) return True
async def test_connect_failure(): """Test connecting to the socket and an exception occurring.""" websocket = Websocket(TEST_API_KEY, TEST_APP_KEY) websocket._sio.connect = AsyncMock(side_effect=SocketIOError()) with pytest.raises(WebsocketError): await websocket.connect()
async def test_connect_sync_success(): """Test connecting to the socket with a sync handler.""" websocket = Websocket(TEST_API_KEY, TEST_APP_KEY) websocket._sio.connect = AsyncMock() websocket._sio.eio._trigger_event = AsyncMock() websocket._sio.namespaces = {"/": 1} async_on_connect = AsyncMock() websocket.async_on_connect(async_on_connect) await websocket.connect() websocket._sio.connect.assert_called_once_with( f"https://rt2.ambientweather.net/?api=1&applicationKey={TEST_APP_KEY}", transports=["websocket"], ) await websocket._sio._trigger_event("connect", "/") async_on_connect.assert_called_once()
async def main() -> None: """Run the websocket example.""" logging.basicConfig(level=logging.INFO) websocket = Websocket(APP_KEY, API_KEY) websocket.on_connect(print_hello) websocket.on_data(print_data) websocket.on_disconnect(print_goodbye) websocket.on_subscribed(print_subscribed) try: await websocket.connect() except WebsocketError as err: _LOGGER.error("There was a websocket error: %s", err) return while True: _LOGGER.info("Simulating some other task occurring...") await asyncio.sleep(5)
async def test_data_async(): """Test data and subscription with async handlers.""" websocket = Websocket(TEST_API_KEY, TEST_APP_KEY) websocket._sio.connect = AsyncMock() websocket._sio.disconnect = AsyncMock() websocket._sio.eio._trigger_event = AsyncMock() websocket._sio.namespaces = {"/": 1} async_on_connect = AsyncMock() async_on_data = AsyncMock() async_on_disconnect = AsyncMock() async_on_subscribed = AsyncMock() websocket.async_on_connect(async_on_connect) websocket.async_on_data(async_on_data) websocket.async_on_disconnect(async_on_disconnect) websocket.async_on_subscribed(async_on_subscribed) await websocket.connect() websocket._sio.connect.assert_called_once_with( f"https://rt2.ambientweather.net/?api=1&applicationKey={TEST_APP_KEY}", transports=["websocket"], ) await websocket._sio._trigger_event("connect", "/") async_on_connect.assert_called_once() await websocket._sio._trigger_event("data", "/", {"foo": "bar"}) async_on_data.assert_called_once() await websocket._sio._trigger_event("subscribed", "/", {"foo": "bar"}) async_on_subscribed.assert_called() await websocket.disconnect() await websocket._sio._trigger_event("disconnect", "/") async_on_disconnect.assert_called_once() websocket._sio.disconnect.assert_called_once()
async def test_reconnect(): """Test that reconnecting to the websocket does the right thing.""" websocket = Websocket(TEST_API_KEY, TEST_APP_KEY) websocket._sio.connect = AsyncMock() websocket._sio.eio._trigger_event = AsyncMock() websocket._sio.namespaces = {"/": 1} async_on_connect = AsyncMock() async_on_disconnect = AsyncMock() websocket.async_on_connect(async_on_connect) websocket.async_on_disconnect(async_on_disconnect) await websocket.reconnect() await websocket._sio._trigger_event("disconnect", "/") async_on_disconnect.assert_called_once() await websocket._sio._trigger_event("connect", "/") async_on_connect.assert_called_once()