2021-07-12 20:09:17 +02:00
|
|
|
import asyncio
|
|
|
|
import socket
|
|
|
|
|
|
|
|
import pytest
|
2021-07-29 16:09:16 +02:00
|
|
|
from mock import AsyncMock, MagicMock, Mock, patch
|
2021-07-12 20:09:17 +02:00
|
|
|
|
2023-01-06 05:24:10 +01:00
|
|
|
from aioesphomeapi._frame_helper import APIPlaintextFrameHelper, Packet
|
2021-07-12 20:09:17 +02:00
|
|
|
from aioesphomeapi.api_pb2 import ConnectResponse, HelloResponse
|
2021-10-05 10:56:35 +02:00
|
|
|
from aioesphomeapi.connection import APIConnection, ConnectionParams, ConnectionState
|
|
|
|
from aioesphomeapi.core import APIConnectionError, RequiresEncryptionAPIError
|
2021-07-12 20:09:17 +02:00
|
|
|
from aioesphomeapi.host_resolver import AddrInfo, IPv4Sockaddr
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.fixture
|
|
|
|
def connection_params() -> ConnectionParams:
|
|
|
|
return ConnectionParams(
|
|
|
|
address="fake.address",
|
|
|
|
port=6052,
|
|
|
|
password=None,
|
|
|
|
client_info="Tests client",
|
|
|
|
keepalive=15.0,
|
|
|
|
zeroconf_instance=None,
|
2021-09-08 23:12:07 +02:00
|
|
|
noise_psk=None,
|
2022-01-20 12:03:36 +01:00
|
|
|
expected_name=None,
|
2021-07-12 20:09:17 +02:00
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.fixture
|
|
|
|
def conn(connection_params) -> APIConnection:
|
|
|
|
async def on_stop():
|
|
|
|
pass
|
|
|
|
|
|
|
|
return APIConnection(connection_params, on_stop)
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.fixture
|
|
|
|
def resolve_host():
|
|
|
|
with patch("aioesphomeapi.host_resolver.async_resolve_host") as func:
|
|
|
|
func.return_value = AddrInfo(
|
|
|
|
family=socket.AF_INET,
|
|
|
|
type=socket.SOCK_STREAM,
|
|
|
|
proto=socket.IPPROTO_TCP,
|
|
|
|
sockaddr=IPv4Sockaddr("10.0.0.512", 6052),
|
|
|
|
)
|
|
|
|
yield func
|
|
|
|
|
|
|
|
|
|
|
|
@pytest.fixture
|
|
|
|
def socket_socket():
|
|
|
|
with patch("socket.socket") as func:
|
|
|
|
yield func
|
|
|
|
|
|
|
|
|
2023-01-06 05:24:10 +01:00
|
|
|
def _get_mock_protocol():
|
|
|
|
def _on_packet(pkt: Packet):
|
|
|
|
pass
|
|
|
|
|
|
|
|
def _on_error(exc: Exception):
|
|
|
|
raise exc
|
|
|
|
|
|
|
|
protocol = APIPlaintextFrameHelper(on_pkt=_on_packet, on_error=_on_error)
|
|
|
|
protocol._connected_event.set()
|
|
|
|
protocol._transport = MagicMock()
|
|
|
|
return protocol
|
|
|
|
|
|
|
|
|
2021-07-12 20:09:17 +02:00
|
|
|
@pytest.mark.asyncio
|
|
|
|
async def test_connect(conn, resolve_host, socket_socket, event_loop):
|
2023-01-06 05:24:10 +01:00
|
|
|
loop = asyncio.get_event_loop()
|
|
|
|
protocol = _get_mock_protocol()
|
|
|
|
with patch.object(event_loop, "sock_connect"), patch.object(
|
|
|
|
loop, "create_connection", return_value=(MagicMock(), protocol)
|
|
|
|
), patch.object(conn, "_connect_start_ping"), patch.object(
|
2021-07-12 20:09:17 +02:00
|
|
|
conn, "send_message_await_response", return_value=HelloResponse()
|
|
|
|
):
|
2021-10-21 19:20:05 +02:00
|
|
|
await conn.connect(login=False)
|
2021-07-12 20:09:17 +02:00
|
|
|
|
|
|
|
assert conn.is_connected
|
2021-10-05 10:56:35 +02:00
|
|
|
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
|
|
async def test_requires_encryption_propagates(conn):
|
2023-01-06 05:24:10 +01:00
|
|
|
loop = asyncio.get_event_loop()
|
|
|
|
protocol = _get_mock_protocol()
|
|
|
|
with patch.object(loop, "create_connection") as create_connection, patch.object(
|
|
|
|
protocol, "perform_handshake"
|
|
|
|
):
|
|
|
|
create_connection.return_value = (MagicMock(), protocol)
|
2021-10-05 10:56:35 +02:00
|
|
|
|
|
|
|
await conn._connect_init_frame_helper()
|
2023-01-06 05:24:10 +01:00
|
|
|
|
2021-10-05 10:56:35 +02:00
|
|
|
with pytest.raises(RequiresEncryptionAPIError):
|
2023-01-06 05:24:10 +01:00
|
|
|
protocol.data_received(b"\x01\x00\x00")
|
2021-10-05 10:56:35 +02:00
|
|
|
await conn._connect_hello()
|