1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285
|
import asyncio
import contextlib
import logging
import socket
from asyncio.streams import StreamReader, StreamWriter
import pytest
from aionut import (
AIONUTClient,
NUTCommandError,
NUTConnectionClosedError,
NUTError,
NUTLoginError,
NUTProtocolError,
NUTShutdownError,
NUTTimeoutError,
)
_LOGGER = logging.getLogger("aionut")
_LOGGER.setLevel(logging.DEBUG)
logging.basicConfig(level=logging.DEBUG)
def test_imports():
assert AIONUTClient()
assert NUTError()
assert NUTProtocolError()
_CLIENTS: set[AIONUTClient] = set()
_SERVERS: set[tuple[asyncio.Server, asyncio.Task[None], socket.socket]] = set()
@pytest.fixture(autouse=True)
async def cleanup():
yield
await asyncio.sleep(0)
_LOGGER.debug("cleanup")
for client in _CLIENTS:
writer = client._writer
assert writer
writer.write_eof()
await writer.drain()
client.shutdown()
await writer.wait_closed()
for server, task, sock in _SERVERS:
server.close()
await server.wait_closed()
sock.close()
for sock in server.sockets:
sock.close()
task.cancel()
with contextlib.suppress(asyncio.CancelledError):
await task
await asyncio.sleep(0)
_CLIENTS.clear()
_SERVERS.clear()
def make_nut_client(
port: int,
username: str | None = "test",
password: str | None = "",
persistent: bool = True,
) -> AIONUTClient:
client = AIONUTClient(
host="localhost",
port=port,
username=username,
password=password,
timeout=0.1,
persistent=persistent,
)
_CLIENTS.add(client)
return client
@pytest.mark.asyncio
async def test_auth_late_auth_failure():
port = await make_fake_nut_server(late_auth_failed=True)
client = make_nut_client(port)
with pytest.raises(NUTLoginError, match="LIST"):
await client.list_ups()
@pytest.mark.asyncio
async def test_no_auth():
port = await make_fake_nut_server()
client = make_nut_client(port, username=None, password=None)
await client.list_ups()
@pytest.mark.asyncio
async def test_auth_bad_username():
port = await make_fake_nut_server(bad_username=True)
client = make_nut_client(port)
with pytest.raises(NUTLoginError, match="USERNAME"):
await client.list_ups()
@pytest.mark.asyncio
async def test_auth_bad_password():
port = await make_fake_nut_server(bad_password=True)
client = make_nut_client(port)
with pytest.raises(NUTLoginError, match="PASSWORD"):
await client.list_ups()
@pytest.mark.asyncio
async def test_list_ups():
port = await make_fake_nut_server()
client = make_nut_client(port)
upses = await client.list_ups()
assert upses == {"test": "bob"}
upses = await client.list_ups()
assert upses == {"test": "bob"}
@pytest.mark.asyncio
async def test_list_ups_no_persist():
port = await make_fake_nut_server()
client = make_nut_client(port, persistent=False)
upses = await client.list_ups()
assert upses == {"test": "bob"}
upses = await client.list_ups()
assert upses == {"test": "bob"}
@pytest.mark.asyncio
async def test_list_vars():
port = await make_fake_nut_server()
client = make_nut_client(port)
vars = await client.list_vars("test")
assert vars == {"x.y": "z"}
@pytest.mark.asyncio
async def test_list_vars_wrong_response():
port = await make_fake_nut_server()
client = make_nut_client(port)
with pytest.raises(NUTProtocolError):
await client.list_vars("wrong")
@pytest.mark.asyncio
async def test_list_command():
port = await make_fake_nut_server()
client = make_nut_client(port)
commands = await client.list_commands("test")
assert commands == {"valid"}
@pytest.mark.asyncio
async def test_list_command_wrong_response():
port = await make_fake_nut_server()
client = make_nut_client(port)
with pytest.raises(NUTProtocolError):
await client.list_commands("wrong")
@pytest.mark.asyncio
async def test_list_ups_first_connection_drop():
port = await make_fake_nut_server(drop_first_connection=True)
client = make_nut_client(port)
upses = await client.list_ups()
assert upses == {"test": "bob"}
@pytest.mark.asyncio
async def test_list_ups_connection_drop():
port = await make_fake_nut_server(drop_connection=True)
client = make_nut_client(port)
with pytest.raises(NUTConnectionClosedError):
await client.list_ups()
@pytest.mark.asyncio
async def test_run_command():
port = await make_fake_nut_server()
client = make_nut_client(port)
with pytest.raises(NUTCommandError, match="UNKNOWN-COMMAND"):
await client.run_command("test", "invalid")
assert await client.run_command("test", "valid") == "OK"
assert await client.run_command("test", "param", "param") == "OK"
@pytest.mark.asyncio
async def test_description():
port = await make_fake_nut_server()
client = make_nut_client(port)
assert await client.description("test") == "demo ups"
@pytest.mark.asyncio
async def test_timeout():
port = await make_fake_nut_server()
client = make_nut_client(port)
with pytest.raises(NUTTimeoutError):
await client.run_command("test", "no_response")
@pytest.mark.asyncio
async def test_use_after_shutdown():
port = await make_fake_nut_server()
client = make_nut_client(port)
client.shutdown()
with pytest.raises(NUTShutdownError):
await client.description("test")
async def make_fake_nut_server(
bad_username: bool = False,
bad_password: bool = False,
late_auth_failed: bool = False,
drop_first_connection: bool = False,
drop_connection: bool = False,
) -> int:
dropped_connection = False
async def handle_client(reader: StreamReader, writer: StreamWriter) -> None:
nonlocal dropped_connection
while True:
if writer.is_closing():
break
command = await reader.readline()
if command == b"":
break
if drop_connection or (drop_first_connection and not dropped_connection):
dropped_connection = True
break
if command.startswith(b"USERNAME"):
if bad_username:
writer.write(b"ERR ACCESS-DENIED\n")
break
writer.write(b"OK\n")
elif command.startswith(b"PASSWORD"):
if bad_password:
writer.write(b"ERR ACCESS-DENIED\n")
break
writer.write(b"OK\n")
elif late_auth_failed:
writer.write(b"ERR ACCESS-DENIED\n")
break
elif command.startswith(b"LIST UPS"):
writer.write(b"BEGIN LIST UPS\n")
writer.write(b'UPS test "bob"\n')
writer.write(b"END LIST UPS\n")
elif command.startswith(b"LIST VAR wrong"):
writer.write(b"OK\n")
elif command.startswith(b"LIST VAR"):
writer.write(b"BEGIN LIST VAR test\n")
writer.write(b'VAR test x.y "z"\n')
writer.write(b"END LIST VAR test\n")
elif command.startswith(b"LIST CMD wrong"):
writer.write(b"OK\n")
elif command.startswith(b"LIST CMD"):
writer.write(b"BEGIN LIST CMD test\n")
writer.write(b'CMD test "valid"\n')
writer.write(b"END LIST CMD test\n")
elif command.startswith(b"INSTCMD test no_response"):
pass
elif command.startswith(b"INSTCMD test invalid"):
writer.write(b"ERR UNKNOWN-COMMAND\n")
elif command.startswith(b"INSTCMD test valid"):
writer.write(b"OK\n")
elif command.startswith(b"INSTCMD test param param"):
writer.write(b"OK\n")
elif command.startswith(b"GET UPSDESC test"):
writer.write(b'UPSDESC test "demo ups"\n')
else:
writer.write(b"ERR\n")
writer.close()
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.bind(("", 0))
sock.listen(1000)
port = sock.getsockname()[1]
server = await asyncio.start_server(handle_client, sock=sock)
task = asyncio.create_task(server.serve_forever())
_SERVERS.add((server, task, sock))
return port
|