forked from svinota/pyroute2
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
dhcp client: add a udhcpd fixture and a test with it
- Loading branch information
Showing
8 changed files
with
384 additions
and
185 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,131 @@ | ||
import abc | ||
import asyncio | ||
from argparse import ArgumentParser | ||
from dataclasses import dataclass | ||
from ipaddress import IPv4Address | ||
from typing import ClassVar, Generic, Literal, Optional, TypeVar | ||
|
||
from ..interfaces import DHCPRangeConfig | ||
|
||
|
||
@dataclass | ||
class DHCPServerConfig: | ||
range: DHCPRangeConfig | ||
interface: str | ||
lease_time: int = 120 # in seconds | ||
max_leases: int = 50 | ||
|
||
|
||
DHCPServerConfigT = TypeVar("DHCPServerConfigT", bound=DHCPServerConfig) | ||
|
||
|
||
class DHCPServerFixture(abc.ABC, Generic[DHCPServerConfigT]): | ||
|
||
BINARY_PATH: ClassVar[Optional[str]] = None | ||
|
||
@classmethod | ||
def get_config_class(cls) -> type[DHCPServerConfigT]: | ||
return cls.__orig_bases__[0].__args__[0] | ||
|
||
def __init__(self, config: DHCPServerConfigT) -> None: | ||
self.config = config | ||
self.stdout: list[str] = [] | ||
self.stderr: list[str] = [] | ||
self.process: Optional[asyncio.subprocess.Process] = None | ||
self.output_poller: Optional[asyncio.Task] = None | ||
|
||
async def _read_output(self, name: Literal['stdout', 'stderr']): | ||
'''Read stdout or stderr until the process exits.''' | ||
stream = getattr(self.process, name) | ||
output = getattr(self, name) | ||
while line := await stream.readline(): | ||
output.append(line.decode().strip()) | ||
|
||
async def _read_outputs(self): | ||
'''Read stdout & stderr until the process exits.''' | ||
assert self.process | ||
await asyncio.gather( | ||
self._read_output('stderr'), self._read_output('stdout') | ||
) | ||
|
||
@abc.abstractmethod | ||
def get_cmdline_options(self) -> tuple[str]: | ||
'''All commandline options passed to the server.''' | ||
|
||
async def __aenter__(self): | ||
'''Start the server process and start polling its output.''' | ||
if not self.BINARY_PATH: | ||
raise RuntimeError( | ||
f"server binary is missing for {type(self.__name__)}" | ||
) | ||
self.process = await asyncio.create_subprocess_exec( | ||
self.BINARY_PATH, | ||
*self.get_cmdline_options(), | ||
stdout=asyncio.subprocess.PIPE, | ||
stderr=asyncio.subprocess.PIPE, | ||
env={'LANG': 'C'}, # usually ensures the output is in english | ||
) | ||
self.output_poller = asyncio.Task(self._read_outputs()) | ||
return self | ||
|
||
async def __aexit__(self, *_): | ||
if self.process: | ||
if self.process.returncode is None: | ||
self.process.terminate() | ||
await self.process.wait() | ||
await self.output_poller | ||
|
||
|
||
def get_psr() -> ArgumentParser: | ||
psr = ArgumentParser() | ||
psr.add_argument('interface', help='Interface to listen on') | ||
psr.add_argument( | ||
'--router', type=IPv4Address, default=None, help='Router IPv4 address.' | ||
) | ||
psr.add_argument( | ||
'--range-start', | ||
type=IPv4Address, | ||
default=IPv4Address('192.168.186.10'), | ||
help='Start of the DHCP client range.', | ||
) | ||
psr.add_argument( | ||
'--range-end', | ||
type=IPv4Address, | ||
default=IPv4Address('192.168.186.100'), | ||
help='End of the DHCP client range.', | ||
) | ||
psr.add_argument( | ||
'--lease-time', | ||
default=120, | ||
type=int, | ||
help='DHCP lease time in seconds (minimum 2 minutes)', | ||
) | ||
psr.add_argument( | ||
'--netmask', type=IPv4Address, default=IPv4Address("255.255.255.0") | ||
) | ||
return psr | ||
|
||
|
||
async def run_fixture_as_main(fixture_cls: type[DHCPServerFixture]): | ||
config_cls = fixture_cls.get_config_class() | ||
args = get_psr().parse_args() | ||
range_config = DHCPRangeConfig( | ||
start=args.range_start, | ||
end=args.range_end, | ||
router=args.router, | ||
netmask=args.netmask, | ||
) | ||
conf = config_cls( | ||
range=range_config, | ||
interface=args.interface, | ||
lease_time=args.lease_time, | ||
) | ||
read_lines: int = 0 | ||
async with fixture_cls(conf) as dhcp_server: | ||
# quick & dirty stderr polling | ||
while True: | ||
if len(dhcp_server.stderr) > read_lines: | ||
read_lines += len(lines := dhcp_server.stderr[read_lines:]) | ||
print(*lines, sep='\n') | ||
else: | ||
await asyncio.sleep(0.2) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,68 @@ | ||
import asyncio | ||
from dataclasses import dataclass | ||
from shutil import which | ||
from typing import AsyncGenerator, ClassVar, Optional | ||
|
||
import pytest | ||
import pytest_asyncio | ||
from fixtures.interfaces import DHCPRangeConfig | ||
|
||
from . import DHCPServerConfig, DHCPServerFixture, run_fixture_as_main | ||
|
||
|
||
@dataclass | ||
class DnsmasqConfig(DHCPServerConfig): | ||
'''Options for the dnsmasq server.''' | ||
|
||
def __iter__(self): | ||
opts = [ | ||
f'--interface={self.interface}', | ||
f'--dhcp-range={self.range.start},' | ||
f'{self.range.end},{self.lease_time}', | ||
f'--dhcp-lease-max={self.max_leases}', | ||
] | ||
if router := self.range.router: | ||
opts.append(f"--dhcp-option=option:router,{router}") | ||
return iter(opts) | ||
|
||
|
||
class DnsmasqFixture(DHCPServerFixture[DnsmasqConfig]): | ||
'''Runs the dnsmasq server as an async context manager.''' | ||
|
||
BINARY_PATH: ClassVar[Optional[str]] = which('dnsmasq') | ||
|
||
def _get_base_cmdline_options(self) -> tuple[str]: | ||
'''The base commandline options for dnsmasq.''' | ||
return ( | ||
'--keep-in-foreground', # self explanatory | ||
'--no-resolv', # don't mess w/ resolv.conf | ||
'--log-facility=-', # log to stdout | ||
'--no-hosts', # don't read /etc/hosts | ||
'--bind-interfaces', # don't bind on wildcard | ||
'--no-ping', # don't ping to check if ips are attributed | ||
) | ||
|
||
def get_cmdline_options(self) -> tuple[str]: | ||
'''All commandline options passed to dnsmasq.''' | ||
return (*self._get_base_cmdline_options(), *self.config) | ||
|
||
|
||
@pytest.fixture | ||
def dnsmasq_config( | ||
veth_pair: tuple[str, str], dhcp_range: DHCPRangeConfig | ||
) -> DnsmasqConfig: | ||
'''dnsmasq options useful for test purposes.''' | ||
return DnsmasqConfig(range=dhcp_range, interface=veth_pair[0]) | ||
|
||
|
||
@pytest_asyncio.fixture | ||
async def dnsmasq( | ||
dnsmasq_config: DnsmasqConfig, | ||
) -> AsyncGenerator[DnsmasqFixture, None]: | ||
'''A dnsmasq instance running for the duration of the test.''' | ||
async with DnsmasqFixture(config=dnsmasq_config) as dnsf: | ||
yield dnsf | ||
|
||
|
||
if __name__ == '__main__': | ||
asyncio.run(run_fixture_as_main(DnsmasqFixture)) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,101 @@ | ||
import asyncio | ||
from dataclasses import dataclass | ||
from pathlib import Path | ||
from shutil import which | ||
from tempfile import TemporaryDirectory | ||
from typing import AsyncGenerator, ClassVar, Optional | ||
|
||
import pytest | ||
import pytest_asyncio | ||
|
||
from ..interfaces import DHCPRangeConfig | ||
from . import DHCPServerConfig, DHCPServerFixture, run_fixture_as_main | ||
|
||
|
||
@dataclass | ||
class UdhcpdConfig(DHCPServerConfig): | ||
arp_ping_timeout_ms: int = 200 # default is 2000 | ||
|
||
|
||
class UdhcpdFixture(DHCPServerFixture[UdhcpdConfig]): | ||
'''Runs the udhcpd server as an async context manager.''' | ||
|
||
BINARY_PATH: ClassVar[Optional[str]] = which('busybox') | ||
|
||
def __init__(self, config): | ||
super().__init__(config) | ||
self._temp_dir: Optional[TemporaryDirectory[str]] = None | ||
|
||
@property | ||
def workdir(self) -> Path: | ||
'''A temporary directory for udhcpd's files.''' | ||
assert self._temp_dir | ||
return Path(self._temp_dir.name) | ||
|
||
@property | ||
def config_file(self) -> Path: | ||
'''The udhcpd config file path.''' | ||
return self.workdir.joinpath("udhcpd.conf") | ||
|
||
async def __aenter__(self): | ||
self._temp_dir = TemporaryDirectory(prefix=type(self).__name__) | ||
self._temp_dir.__enter__() | ||
self.config_file.write_text(self.generate_config()) | ||
return await super().__aenter__() | ||
|
||
def generate_config(self) -> str: | ||
'''Generate the contents of udhcpd's config file.''' | ||
cfg = self.config | ||
base_workfile = self.workdir.joinpath(self.config.interface) | ||
lease_file = base_workfile.with_suffix(".leases") | ||
pidfile = base_workfile.with_suffix(".pid") | ||
lines = [ | ||
("start", cfg.range.start), | ||
("end", cfg.range.end), | ||
("max_leases", cfg.max_leases), | ||
("interface", cfg.interface), | ||
("lease_file", lease_file), | ||
("pidfile", pidfile), | ||
("opt lease", cfg.lease_time), | ||
("opt router", cfg.range.router), | ||
] | ||
return "\n".join(f"{opt}\t{value}" for opt, value in lines) | ||
|
||
async def __aexit__(self, *_): | ||
await super().__aexit__(*_) | ||
self._temp_dir.__exit__(*_) | ||
|
||
def get_cmdline_options(self) -> tuple[str]: | ||
'''All commandline options passed to udhcpd.''' | ||
return ( | ||
'udhcpd', | ||
'-f', # run in foreground | ||
'-a', | ||
str(self.config.arp_ping_timeout_ms), | ||
str(self.config_file), | ||
) | ||
|
||
|
||
@pytest.fixture | ||
def udhcpd_config( | ||
veth_pair: tuple[str, str], dhcp_range: DHCPRangeConfig | ||
) -> UdhcpdConfig: | ||
'''udhcpd options useful for test purposes.''' | ||
return UdhcpdConfig( | ||
range=dhcp_range, | ||
interface=veth_pair[0], | ||
lease_time=1, # very short leases for tests | ||
) | ||
|
||
|
||
@pytest_asyncio.fixture | ||
async def udhcpd( | ||
udhcpd_config: UdhcpdConfig, | ||
) -> AsyncGenerator[UdhcpdFixture, None]: | ||
'''An udhcpd instance running for the duration of the test.''' | ||
async with UdhcpdFixture(config=udhcpd_config) as dhcp_server: | ||
yield dhcp_server | ||
|
||
|
||
if __name__ == '__main__': | ||
asyncio.run(run_fixture_as_main(UdhcpdFixture)) |
Oops, something went wrong.