Add additional coverage for starting noise connections (#653)

This commit is contained in:
J. Nick Koston 2023-11-21 16:42:23 +01:00 committed by GitHub
parent 0afa8c6832
commit 2501ea4275
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 50 additions and 4 deletions

View File

@ -20,7 +20,8 @@ from google.protobuf import message
import aioesphomeapi.host_resolver as hr import aioesphomeapi.host_resolver as hr
from ._frame_helper import APINoiseFrameHelper, APIPlaintextFrameHelper from ._frame_helper.noise import APINoiseFrameHelper
from ._frame_helper.plain_text import APIPlaintextFrameHelper
from .api_pb2 import ( # type: ignore from .api_pb2 import ( # type: ignore
ConnectRequest, ConnectRequest,
ConnectResponse, ConnectResponse,

View File

@ -57,13 +57,33 @@ def connection_params() -> ConnectionParams:
@pytest.fixture @pytest.fixture
def conn(connection_params) -> APIConnection: def noise_connection_params() -> ConnectionParams:
async def on_stop(expected_disconnect: bool) -> None: return ConnectionParams(
pass address="fake.address",
port=6052,
password=None,
client_info="Tests client",
keepalive=KEEP_ALIVE_INTERVAL,
zeroconf_manager=ZeroconfManager(),
noise_psk="QRTIErOb/fcE9Ukd/5qA3RGYMn0Y+p06U58SCtOXvPc=",
expected_name="test",
)
async def on_stop(expected_disconnect: bool) -> None:
pass
@pytest.fixture
def conn(connection_params: ConnectionParams) -> APIConnection:
return APIConnection(connection_params, on_stop) return APIConnection(connection_params, on_stop)
@pytest.fixture
def noise_conn(noise_connection_params: ConnectionParams) -> APIConnection:
return APIConnection(noise_connection_params, on_stop)
@pytest_asyncio.fixture(name="plaintext_connect_task_no_login") @pytest_asyncio.fixture(name="plaintext_connect_task_no_login")
async def plaintext_connect_task_no_login( async def plaintext_connect_task_no_login(
conn: APIConnection, resolve_host, socket_socket, event_loop conn: APIConnection, resolve_host, socket_socket, event_loop

View File

@ -596,6 +596,31 @@ async def test_init_plaintext_with_wrong_preamble(conn: APIConnection):
await task await task
@pytest.mark.asyncio
async def test_init_noise_with_wrong_byte_marker(noise_conn: APIConnection) -> None:
loop = asyncio.get_event_loop()
transport = MagicMock()
protocol: APINoiseFrameHelper | None = None
async def _create_connection(create, sock, *args, **kwargs):
nonlocal protocol
protocol = create()
protocol.connection_made(transport)
return transport, protocol
with patch.object(loop, "create_connection", side_effect=_create_connection):
task = asyncio.create_task(noise_conn._connect_init_frame_helper())
await asyncio.sleep(0)
assert protocol is not None
assert isinstance(noise_conn._frame_helper, APINoiseFrameHelper)
protocol.data_received(b"\x00\x00\x00")
with pytest.raises(ProtocolAPIError, match="Marker byte invalid"):
await task
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_eof_received_closes_connection( async def test_eof_received_closes_connection(
plaintext_connect_task_with_login: tuple[ plaintext_connect_task_with_login: tuple[