test_base_client.py

#  Copyright 2019-2021, 2024  Dom Sekotill <dom.sekotill@kodo.org.uk>
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""
Test cases for wpa_supplicant.client.base.BaseClient
"""

import unittest
from unittest.mock import AsyncMock

import anyio

from tests._anyio import patch_connect
from tests._anyio import patch_send
from wpa_supplicant import errors
from wpa_supplicant.client import base


class ConnectTests(unittest.IsolatedAsyncioTestCase):
	"""
	Tests for the connect() method
	"""

	async def test_connect(self) -> None:
		"""
		Check connect() calls socket.connect()
		"""
		with patch_connect() as connect_mock, patch_send():
			async with base.BaseClient() as client:
				await client.connect("foo")

		connect_mock.assert_awaited_once_with("foo")

	async def test_connect_timeout_1(self) -> None:
		"""
		Check a socket.connect() delay causes TimeoutError to be raised
		"""
		with patch_connect(2.0), patch_send():
			async with base.BaseClient() as client:
				with self.assertRaises(TimeoutError):
					await client.connect("foo")

	async def test_connect_timeout_2(self) -> None:
		"""
		Check a send/recv delay causes a TimeoutError to be raised
		"""
		with patch_connect(), patch_send(2.0):
			async with base.BaseClient() as client:
				with self.assertRaises(TimeoutError):
					await client.connect("foo")


class SendMessageTests(unittest.IsolatedAsyncioTestCase):
	"""
	Tests for the send_command() method
	"""

	def setUp(self) -> None:
		self.client = client = base.BaseClient()
		client.sock = AsyncMock(spec=anyio.abc.SocketStream)
		client.sock.send.return_value = None
		assert isinstance(client.sock, anyio.abc.SocketStream)

	async def test_simple(self) -> None:
		"""
		Check that a response is processed after a command
		"""
		async with self.client as client:
			client.sock.receive.return_value = b"OK"
			assert await client.send_command("SOME_COMMAND") is None

	async def test_simple_expect(self) -> None:
		"""
		Check that an alternate expected response is processed
		"""
		async with self.client as client:
			client.sock.receive.return_value = b"PONG"
			assert await client.send_command("PING", expect="PONG") is None

	async def test_simple_no_expect(self) -> None:
		"""
		Check that an unexpected response raises an UnexpectedResponseError
		"""
		async with self.client as client:
			client.sock.receive.return_value = b"DING"
			with self.assertRaises(errors.UnexpectedResponseError):
				await client.send_command("PING")
			with self.assertRaises(errors.UnexpectedResponseError):
				await client.send_command("PING", expect="PONG")

	async def test_simple_convert(self) -> None:
		"""
		Check that a response is passed through a converter if given
		"""
		async with self.client as client:
			client.sock.receive.return_value = b"FOO\nBAR\nBAZ\n"
			self.assertListEqual(
				await client.send_command(
					"SOME_COMMAND", convert=lambda x: x.splitlines(),
				),
				["FOO", "BAR", "BAZ"],
			)

	async def test_simple_convert_over_expect(self) -> None:
		"""
		Check that 'convert' overrides 'expect'
		"""
		async with self.client as client:
			client.sock.receive.return_value = b"FOO\nBAR\nBAZ\n"
			self.assertListEqual(
				await client.send_command(
					"SOME_COMMAND", convert=lambda x: x.splitlines(), expect="PONG",
				),
				["FOO", "BAR", "BAZ"],
			)

	async def test_simple_fail(self) -> None:
		"""
		Check that a response of 'FAIL' causes CommandFailed to be raised
		"""
		async with self.client as client:
			client.sock.receive.return_value = b"FAIL"
			with self.assertRaises(errors.CommandFailed):
				await client.send_command("SOME_COMMAND")

	async def test_simple_bad_command(self) -> None:
		"""
		Check that a response of 'UNKNOWN COMMAND' causes ValueError to be raised
		"""
		async with self.client as client:
			client.sock.receive.return_value = b"UNKNOWN COMMAND"
			with self.assertRaises(ValueError):
				await client.send_command("SOME_COMMAND")

	async def test_interleaved(self) -> None:
		"""
		Check that messages are processed alongside replies
		"""
		async with self.client as client:
			client.sock.receive.side_effect = [
				b"<2>SOME-MESSAGE",
				b"<1>SOME-OTHER-MESSAGE with|args",
				b"OK",
				b"<2>SOME-MESSAGE",
			]
			assert await client.send_command("SOME_COMMAND") is None

	async def test_unexpected(self) -> None:
		"""
		Check that unexpected replies are logged cleanly
		"""
		async with self.client as client:
			client.sock.receive.side_effect = [
				b"OK",  # Response to "ATTACH"
				b"UNEXPECTED1",
				b"UNEXPECTED2",
				b"<2>CTRL-EVENT-EXAMPLE",
				b"OK",  # Response to "DETACH"
			]
			assert await client.event("CTRL-EVENT-EXAMPLE")

	async def test_unconnected(self) -> None:
		"""
		Check that calling send_command() on an unconnected client raises RuntimeError
		"""
		client = base.BaseClient()

		with self.assertRaises(RuntimeError):
			await client.send_command("SOME_COMMAND")

	async def test_multi_task(self) -> None:
		"""
		Check that calling send_command() from multiple tasks works as expected
		"""
		recv_responses = iter([
			(0.0, b"OK"),           # Response to ATTACH
			(0.5, b"OK"),           # Response to SOME_COMMAND1
			(0.2, b"<2>CTRL-FOO"),  # Event
			(0.1, b"REPLY2"),       # Response to SOME_COMMAND2
			(0.0, b"OK"),           # Response to DETACH
		])

		async def recv() -> bytes:
			delay, data = next(recv_responses)
			await anyio.sleep(delay)
			return data

		async with self.client as client, anyio.create_task_group() as task_group:
			client.sock.receive.side_effect = recv

			@task_group.start_soon
			async def wait_for_event() -> None:
				self.assertTupleEqual(
					await client.event("CTRL-FOO"),
					(base.EventPriority.INFO, "CTRL-FOO", None),
				)
			await anyio.sleep(0.1)  # Ensure send_command("ATTACH") has been sent

			task_group.start_soon(client.send_command, "SOME_COMMAND1")
			await anyio.sleep(0.1)  # Ensure send_command("SOME_COMMAND1") has been sent

			# At this point the response to SOME_COMMAND1 is still delayed
			await client.send_command("SOME_COMMAND2", expect="REPLY2")

	async def test_multi_task_decode_error(self) -> None:
		"""
		Check that decode errors closes the socket and causes all tasks to raise EOFError
		"""
		recv_responses = [
			b"OK",           # Response to ATTACH
			b"\xa5\x8b",     # Undecodable input
			anyio.EndOfStream,
			anyio.EndOfStream,
		]

		async with self.client as client, anyio.create_task_group() as task_group:
			client.sock.receive.side_effect = recv_responses

			@task_group.start_soon
			async def wait_for_event() -> None:
				with self.assertRaises(anyio.ClosedResourceError):
					await client.event("CTRL-FOO")
			await anyio.sleep(0.1)  # Ensure send_command("ATTACH") has been sent

			with self.assertRaises(anyio.ClosedResourceError):
				await client.send_command("SOME_COMMAND", expect="REPLY")


class EventTests(unittest.IsolatedAsyncioTestCase):
	"""
	Tests for the event() method
	"""

	def setUp(self) -> None:
		self.client = client = base.BaseClient()
		client.sock = AsyncMock()
		client.sock.send.return_value = None

	async def test_simple(self) -> None:
		"""
		Check that an awaited message is returned when is arrives
		"""
		with anyio.fail_after(2):
			async with self.client as client:
				client.sock.receive.side_effect = [
					b"OK",  # Respond to ATTACH
					b"<2>CTRL-EVENT-EXAMPLE",
					b"OK",  # Respond to DETACH
				]
				prio, evt, args = await client.event("CTRL-EVENT-EXAMPLE")
				assert prio == 2
				assert evt == "CTRL-EVENT-EXAMPLE"
				assert args is None

	async def test_multiple(self) -> None:
		"""
		Check that an awaited messages is returned when it arrives between others
		"""
		with anyio.fail_after(2):
			async with self.client as client:
				client.sock.receive.side_effect = [
					b"OK",  # Respond to ATTACH
					b"<1>OTHER-MESSAGE",
					b"<2>CTRL-EVENT-OTHER",
					b"<4>CTRL-EVENT-EXAMPLE",
					b"OK",  # Respond to DETACH
					b"<3>OTHER-MESSAGE",
				]
				prio, evt, args = await client.event("CTRL-EVENT-EXAMPLE")
				assert prio == 4
				assert evt == "CTRL-EVENT-EXAMPLE"
				assert args is None

	async def test_wait_multiple(self) -> None:
		"""
		Check that the first of several awaited events is returned
		"""
		with anyio.fail_after(2):
			async with self.client as client:
				client.sock.receive.side_effect = [
					b"OK",  # Respond to ATTACH
					b"<1>OTHER-MESSAGE",
					b"<2>CTRL-EVENT-OTHER",
					b"<4>CTRL-EVENT-EXAMPLE3",
					b"<4>CTRL-EVENT-EXAMPLE1",
					b"OK",  # Respond to DETACH
					b"<3>CTRL-EVENT-OTHER",
				]
				prio, evt, args = await client.event(
					"CTRL-EVENT-EXAMPLE1", "CTRL-EVENT-EXAMPLE2", "CTRL-EVENT-EXAMPLE3",
				)
				assert prio == 4
				assert evt == "CTRL-EVENT-EXAMPLE3"
				assert args is None

	async def test_interleaved(self) -> None:
		"""
		Check that messages are processed as well as replies
		"""
		with anyio.fail_after(2):
			async with self.client as client:
				client.sock.receive.side_effect = [
					b"<1>OTHER-MESSAGE",
					b"OK",  # Respond to SOME_COMMAND
					b"OK",  # Respond to ATTACH
					b"<2>CTRL-EVENT-OTHER",
					b"<4>CTRL-EVENT-EXAMPLE",
					b"<3>CTRL-EVENT-OTHER",
					b"OK",  # Respond to DETACH
					b"FOO",
				]

				assert await client.send_command("SOME_COMMAND") is None

				prio, evt, args = await client.event("CTRL-EVENT-EXAMPLE")
				assert prio == 4
				assert evt == "CTRL-EVENT-EXAMPLE"
				assert args is None

				assert await client.send_command("SOME_COMMAND", expect="FOO") is None

	async def test_unconnected(self) -> None:
		"""
		Check that calling event() on an unconnected client raises RuntimeError
		"""
		client = base.BaseClient()

		with self.assertRaises(RuntimeError):
			await client.event("some", "events")