aioesphomeapi/tests/test_connection.py
2021-10-21 19:20:05 +02:00

79 lines
2.2 KiB
Python

import asyncio
import socket
import pytest
from mock import AsyncMock, MagicMock, Mock, patch
from aioesphomeapi.api_pb2 import ConnectResponse, HelloResponse
from aioesphomeapi.connection import APIConnection, ConnectionParams, ConnectionState
from aioesphomeapi.core import APIConnectionError, RequiresEncryptionAPIError
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,
noise_psk=None,
)
@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
@pytest.mark.asyncio
async def test_connect(conn, resolve_host, socket_socket, event_loop):
with patch.object(event_loop, "sock_connect"), patch(
"asyncio.open_connection", return_value=(None, None)
), patch.object(conn, "_read_loop"), patch.object(
conn, "_connect_start_ping"
), patch.object(
conn, "send_message_await_response", return_value=HelloResponse()
):
await conn.connect(login=False)
assert conn.is_connected
@pytest.mark.asyncio
async def test_requires_encryption_propagates(conn):
with patch("asyncio.open_connection") as openc:
reader = MagicMock()
writer = MagicMock()
openc.return_value = (reader, writer)
writer.drain = AsyncMock()
reader.readexactly = AsyncMock()
reader.readexactly.return_value = b"\x01"
await conn._connect_init_frame_helper()
with pytest.raises(RequiresEncryptionAPIError):
await conn._connect_hello()