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()