Skip to content

Commit

Permalink
dhcp client: add a udhcpd fixture and a test with it
Browse files Browse the repository at this point in the history
  • Loading branch information
etene committed Jan 17, 2025
1 parent 1e195c4 commit d6fa1e5
Show file tree
Hide file tree
Showing 8 changed files with 384 additions and 185 deletions.
12 changes: 10 additions & 2 deletions pyroute2/dhcp/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,13 +59,21 @@ def __init__(

# "public api"

async def wait_for_state(self, state: Optional[fsm.State]) -> None:
async def wait_for_state(
self, state: Optional[fsm.State], timeout: Optional[float] = None
) -> None:
'''Waits until the client is in the target state.
Since the state is set to None upon exit,
you can also pass None to wait for the client to stop.
'''
await self._states[state].wait()
try:
await asyncio.wait_for(self._states[state].wait(), timeout=timeout)
except TimeoutError as err:
raise TimeoutError(
f"Timed out waiting for the {state} state. "
f"Current state: {self.state}"
) from err

@fsm.state_guard(fsm.State.INIT, fsm.State.INIT_REBOOT)
async def bootstrap(self):
Expand Down
3 changes: 2 additions & 1 deletion tests/test_linux/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,8 @@
from uuid import uuid4

import pytest
from fixtures.dnsmasq import dnsmasq, dnsmasq_options # noqa: F401
from fixtures.dhcp_servers.dnsmasq import dnsmasq, dnsmasq_config # noqa: F401
from fixtures.dhcp_servers.udhcpd import udhcpd, udhcpd_config # noqa: F401
from fixtures.interfaces import dhcp_range, veth_pair # noqa: F401
from pr2test.context_manager import NDBContextManager, SpecContextManager
from utils import require_user
Expand Down
131 changes: 131 additions & 0 deletions tests/test_linux/fixtures/dhcp_servers/__init__.py
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)
68 changes: 68 additions & 0 deletions tests/test_linux/fixtures/dhcp_servers/dnsmasq.py
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))
101 changes: 101 additions & 0 deletions tests/test_linux/fixtures/dhcp_servers/udhcpd.py
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))
Loading

0 comments on commit d6fa1e5

Please sign in to comment.