From ae2a5cc8e8e1f801899fb4b69ce9ffdb7aa64e9f Mon Sep 17 00:00:00 2001 From: Yanks Yoon <37652070+yanksyoon@users.noreply.github.com> Date: Thu, 25 Jan 2024 22:03:51 +0800 Subject: [PATCH] feat: support random choice from tmate relations (#208) * feat: support random choice from tmate relations * chore: run src-docs * chore: type hint use default list * test: add failed test description * fix: json parsing --- src-docs/charm_state.py.md | 6 ++-- src-docs/runner.py.md | 6 ++-- src-docs/runner_type.py.md | 2 +- src/charm_state.py | 56 +++++++++++++++++-------------- src/runner.py | 11 ++++-- src/runner_manager.py | 4 +-- src/runner_type.py | 4 +-- tests/unit/factories.py | 48 ++++++++++++++++++++++++++ tests/unit/requirements.txt | 1 + tests/unit/test_charm_state.py | 12 +++---- tests/unit/test_runner.py | 46 ++++++++++++++++++++++++- tests/unit/test_runner_manager.py | 2 +- tox.ini | 1 + 13 files changed, 152 insertions(+), 47 deletions(-) create mode 100644 tests/unit/factories.py create mode 100644 tests/unit/requirements.txt diff --git a/src-docs/charm_state.py.md b/src-docs/charm_state.py.md index 2c3119bdd..f4f3c00ee 100644 --- a/src-docs/charm_state.py.md +++ b/src-docs/charm_state.py.md @@ -183,7 +183,7 @@ SSH connection information for debug workflow. ### classmethod `from_charm` ```python -from_charm(charm: CharmBase) → Optional[ForwardRef('SSHDebugInfo')] +from_charm(charm: CharmBase) → list['SSHDebugInfo'] ``` Initialize the SSHDebugInfo from charm relation data. @@ -208,14 +208,14 @@ The charm state. - `proxy_config`: Proxy-related configuration. - `charm_config`: Configuration of the juju charm. - `arch`: The underlying compute architecture, i.e. x86_64, amd64, arm64/aarch64. - - `ssh_debug_info`: The SSH debug connection configuration information. + - `ssh_debug_infos`: SSH debug connections configuration information. --- - + ### classmethod `from_charm` diff --git a/src-docs/runner.py.md b/src-docs/runner.py.md index 3de69178d..9916319b7 100644 --- a/src-docs/runner.py.md +++ b/src-docs/runner.py.md @@ -39,7 +39,7 @@ The configuration values for creating a single runner instance. ## class `Runner` Single instance of GitHub self-hosted runner. - + ### function `__init__` @@ -67,7 +67,7 @@ Construct the runner instance. --- - + ### function `create` @@ -91,7 +91,7 @@ Create the runner instance on LXD and register it on GitHub. --- - + ### function `remove` diff --git a/src-docs/runner_type.py.md b/src-docs/runner_type.py.md index a87aea240..5bf0f4743 100644 --- a/src-docs/runner_type.py.md +++ b/src-docs/runner_type.py.md @@ -83,7 +83,7 @@ Configuration for runner. - `path`: GitHub repository path in the format '/', or the GitHub organization name. - `proxies`: HTTP(S) proxy settings. - `dockerhub_mirror`: URL of dockerhub mirror to use. - - `ssh_debug_info`: The SSH debug server connection metadata. + - `ssh_debug_infos`: The SSH debug server connections metadata. diff --git a/src/charm_state.py b/src/charm_state.py index d51a9b71f..21bcd1475 100644 --- a/src/charm_state.py +++ b/src/charm_state.py @@ -200,31 +200,37 @@ class SSHDebugInfo(BaseModel): ed25519_fingerprint: str = Field(pattern="^SHA256:.*") @classmethod - def from_charm(cls, charm: CharmBase) -> Optional["SSHDebugInfo"]: + def from_charm(cls, charm: CharmBase) -> list["SSHDebugInfo"]: """Initialize the SSHDebugInfo from charm relation data. Args: charm: The charm instance. """ + ssh_debug_connections: list[SSHDebugInfo] = [] relations = charm.model.relations[DEBUG_SSH_INTEGRATION_NAME] if not relations or not (relation := relations[0]).units: - return None - target_unit = next(iter(relation.units)) - relation_data = relation.data[target_unit] - if ( - not (host := relation_data.get("host")) - or not (port := relation_data.get("port")) - or not (rsa_fingerprint := relation_data.get("rsa_fingerprint")) - or not (ed25519_fingerprint := relation_data.get("ed25519_fingerprint")) - ): - logger.warning("%s relation data not yet ready.", DEBUG_SSH_INTEGRATION_NAME) - return None - return SSHDebugInfo( - host=host, - port=port, - rsa_fingerprint=rsa_fingerprint, - ed25519_fingerprint=ed25519_fingerprint, - ) + return ssh_debug_connections + for unit in relation.units: + relation_data = relation.data[unit] + if ( + not (host := relation_data.get("host")) + or not (port := relation_data.get("port")) + or not (rsa_fingerprint := relation_data.get("rsa_fingerprint")) + or not (ed25519_fingerprint := relation_data.get("ed25519_fingerprint")) + ): + logger.warning( + "%s relation data for %s not yet ready.", DEBUG_SSH_INTEGRATION_NAME, unit.name + ) + continue + ssh_debug_connections.append( + SSHDebugInfo( + host=host, + port=port, + rsa_fingerprint=rsa_fingerprint, + ed25519_fingerprint=ed25519_fingerprint, + ) + ) + return ssh_debug_connections @dataclasses.dataclass(frozen=True) @@ -236,14 +242,14 @@ class State: proxy_config: Proxy-related configuration. charm_config: Configuration of the juju charm. arch: The underlying compute architecture, i.e. x86_64, amd64, arm64/aarch64. - ssh_debug_info: The SSH debug connection configuration information. + ssh_debug_infos: SSH debug connections configuration information. """ is_metrics_logging_available: bool proxy_config: ProxyConfig charm_config: CharmConfig arch: ARCH - ssh_debug_info: Optional[SSHDebugInfo] + ssh_debug_infos: list[SSHDebugInfo] @classmethod def from_charm(cls, charm: CharmBase) -> "State": @@ -292,7 +298,7 @@ def from_charm(cls, charm: CharmBase) -> "State": raise CharmConfigInvalidError(f"Unsupported architecture {exc.arch}") from exc try: - ssh_debug_info = SSHDebugInfo.from_charm(charm) + ssh_debug_infos = SSHDebugInfo.from_charm(charm) except ValidationError as exc: logger.error("Invalid SSH debug info: %s.", exc) raise CharmConfigInvalidError("Invalid SSH Debug info") from exc @@ -302,16 +308,16 @@ def from_charm(cls, charm: CharmBase) -> "State": proxy_config=proxy_config, charm_config=charm_config, arch=arch, - ssh_debug_info=ssh_debug_info, + ssh_debug_infos=ssh_debug_infos, ) state_dict = dataclasses.asdict(state) # Convert pydantic object to python object serializable by json module. state_dict["proxy_config"] = json.loads(state_dict["proxy_config"].json()) state_dict["charm_config"] = json.loads(state_dict["charm_config"].json()) - state_dict["ssh_debug_info"] = ( - json.loads(state_dict["ssh_debug_info"].json()) if ssh_debug_info else None - ) + state_dict["ssh_debug_infos"] = [ + debug_info.json() for debug_info in state_dict["ssh_debug_infos"] + ] json_data = json.dumps(state_dict, ensure_ascii=False) CHARM_STATE_PATH.write_text(json_data, encoding="utf-8") diff --git a/src/runner.py b/src/runner.py index 5fa84deab..2bfbb6e45 100644 --- a/src/runner.py +++ b/src/runner.py @@ -13,6 +13,7 @@ import json import logging import pathlib +import secrets import textwrap import time from dataclasses import dataclass @@ -22,7 +23,7 @@ import yaml import shared_fs -from charm_state import ARCH +from charm_state import ARCH, SSHDebugInfo from errors import ( CreateSharedFilesystemError, LxdError, @@ -671,9 +672,13 @@ def _configure_runner(self) -> None: # As the user already has sudo access, this does not give the user any additional access. self.instance.execute(["/usr/bin/sudo", "chmod", "777", "/usr/local/bin"]) + selected_ssh_connection: SSHDebugInfo | None = ( + secrets.choice(self.config.ssh_debug_infos) if self.config.ssh_debug_infos else None + ) + logger.info("SSH Debug info: %s", selected_ssh_connection) # Load `/etc/environment` file. environment_contents = self._clients.jinja.get_template("environment.j2").render( - proxies=self.config.proxies, ssh_debug_info=self.config.ssh_debug_info + proxies=self.config.proxies, ssh_debug_info=selected_ssh_connection ) self._put_file("/etc/environment", environment_contents) @@ -682,7 +687,7 @@ def _configure_runner(self) -> None: proxies=self.config.proxies, pre_job_script=str(self.pre_job_script), dockerhub_mirror=self.config.dockerhub_mirror, - ssh_debug_info=self.config.ssh_debug_info, + ssh_debug_info=selected_ssh_connection, ) self._put_file(str(self.env_file), env_contents) self.instance.execute(["/usr/bin/chown", "ubuntu:ubuntu", str(self.env_file)]) diff --git a/src/runner_manager.py b/src/runner_manager.py index 187d1590e..d85203d63 100644 --- a/src/runner_manager.py +++ b/src/runner_manager.py @@ -409,7 +409,7 @@ def _spawn_new_runners(self, count: int, resources: VirtualMachineResources): path=self.config.path, proxies=self.proxies, name=self._generate_runner_name(), - ssh_debug_info=self.config.charm_state.ssh_debug_info, + ssh_debug_infos=self.config.charm_state.ssh_debug_infos, ) runner = Runner(self._clients, config, RunnerStatus()) try: @@ -611,7 +611,7 @@ def create_runner_info( name=name, path=self.config.path, proxies=self.proxies, - ssh_debug_info=self.config.charm_state.ssh_debug_info, + ssh_debug_infos=self.config.charm_state.ssh_debug_infos, ) return Runner( self._clients, diff --git a/src/runner_type.py b/src/runner_type.py index 6afed8766..abdb36ae2 100644 --- a/src/runner_type.py +++ b/src/runner_type.py @@ -77,7 +77,7 @@ class RunnerConfig: # pylint: disable=too-many-instance-attributes name. proxies: HTTP(S) proxy settings. dockerhub_mirror: URL of dockerhub mirror to use. - ssh_debug_info: The SSH debug server connection metadata. + ssh_debug_infos: The SSH debug server connections metadata. """ app_name: str @@ -87,7 +87,7 @@ class RunnerConfig: # pylint: disable=too-many-instance-attributes path: GithubPath proxies: ProxySetting dockerhub_mirror: str | None = None - ssh_debug_info: SSHDebugInfo | None = None + ssh_debug_infos: SSHDebugInfo | None = None @dataclass diff --git a/tests/unit/factories.py b/tests/unit/factories.py new file mode 100644 index 000000000..55733a2e3 --- /dev/null +++ b/tests/unit/factories.py @@ -0,0 +1,48 @@ +# Copyright 2024 Canonical Ltd. +# See LICENSE file for licensing details. + +"""Factories for generating test data.""" + +# The factory definitions don't need public methods +# pylint: disable=too-few-public-methods + +import random +from typing import Generic, TypeVar + +import factory +import factory.fuzzy +from pydantic.networks import IPvAnyAddress + +from charm_state import SSHDebugInfo + +T = TypeVar("T") + + +class BaseMetaFactory(Generic[T], factory.base.FactoryMetaClass): + """Used for type hints of factories.""" + + # No need for docstring because it is used for type hints + def __call__(cls, *args, **kwargs) -> T: # noqa: N805 + """Used for type hints of factories.""" # noqa: DCO020 + return super().__call__(*args, **kwargs) # noqa: DCO030 + + +# The attributes of these classes are generators for the attributes of the meta class +# mypy incorrectly believes the factories don't support metaclass +class SSHDebugInfoFactory( + factory.Factory, metaclass=BaseMetaFactory[SSHDebugInfo] # type: ignore +): + # Docstrings have been abbreviated for factories, checking for docstrings on model attributes + # can be skipped. + """Generate PathInfos.""" # noqa: DCO060 + + class Meta: + """Configuration for factory.""" # noqa: DCO060 + + model = SSHDebugInfo + abstract = False + + host: IPvAnyAddress = factory.Faker("ipv4") + port: int = factory.LazyAttribute(lambda n: random.randrange(1024, 65536)) + rsa_fingerprint: str = factory.fuzzy.FuzzyText(prefix="SHA256:") + ed25519_fingerprint: str = factory.fuzzy.FuzzyText(prefix="SHA256:") diff --git a/tests/unit/requirements.txt b/tests/unit/requirements.txt new file mode 100644 index 000000000..c9c3f2669 --- /dev/null +++ b/tests/unit/requirements.txt @@ -0,0 +1 @@ +factory-boy>=3,<4 diff --git a/tests/unit/test_charm_state.py b/tests/unit/test_charm_state.py index c3e147e55..8e1c77ee6 100644 --- a/tests/unit/test_charm_state.py +++ b/tests/unit/test_charm_state.py @@ -132,7 +132,7 @@ def test_ssh_debug_info_from_charm_no_relations(): mock_charm = MagicMock(spec=GithubRunnerCharm) mock_charm.model.relations = {DEBUG_SSH_INTEGRATION_NAME: []} - assert SSHDebugInfo.from_charm(mock_charm) is None + assert not SSHDebugInfo.from_charm(mock_charm) @pytest.mark.parametrize( @@ -220,11 +220,11 @@ def test_from_charm_ssh_debug_info(): mock_charm.app.name = "github-runner-operator" mock_charm.unit.name = "github-runner-operator/0" - ssh_debug_info = State.from_charm(mock_charm).ssh_debug_info - assert str(ssh_debug_info.host) == mock_relation_data["host"] - assert str(ssh_debug_info.port) == mock_relation_data["port"] - assert ssh_debug_info.rsa_fingerprint == mock_relation_data["rsa_fingerprint"] - assert ssh_debug_info.ed25519_fingerprint == mock_relation_data["ed25519_fingerprint"] + ssh_debug_infos = State.from_charm(mock_charm).ssh_debug_infos + assert str(ssh_debug_infos[0].host) == mock_relation_data["host"] + assert str(ssh_debug_infos[0].port) == mock_relation_data["port"] + assert ssh_debug_infos[0].rsa_fingerprint == mock_relation_data["rsa_fingerprint"] + assert ssh_debug_infos[0].ed25519_fingerprint == mock_relation_data["ed25519_fingerprint"] def test_invalid_runner_storage(): diff --git a/tests/unit/test_runner.py b/tests/unit/test_runner.py index 889b8f787..3ecfe49de 100644 --- a/tests/unit/test_runner.py +++ b/tests/unit/test_runner.py @@ -12,11 +12,13 @@ import pytest from _pytest.monkeypatch import MonkeyPatch +from charm_state import SSHDebugInfo from errors import CreateSharedFilesystemError, RunnerCreateError, RunnerRemoveError from runner import CreateRunnerConfig, Runner, RunnerConfig, RunnerStatus from runner_manager_type import RunnerManagerClients from runner_type import GithubOrg, GithubRepo, VirtualMachineResources from shared_fs import SharedFilesystem +from tests.unit.factories import SSHDebugInfoFactory from tests.unit.mock import ( MockLxdClient, MockRepoPolicyComplianceClient, @@ -89,6 +91,12 @@ def jinja2_environment_fixture() -> MagicMock: return jinja2_mock +@pytest.fixture(scope="function", name="ssh_debug_infos") +def ssh_debug_infos_fixture() -> list[SSHDebugInfo]: + """A list of randomly generated ssh_debug_infos.""" + return SSHDebugInfoFactory.create_batch(size=100) + + @pytest.fixture( scope="function", name="runner", @@ -100,7 +108,13 @@ def jinja2_environment_fixture() -> MagicMock: ), ], ) -def runner_fixture(request, lxd: MockLxdClient, jinja: MagicMock, tmp_path: Path): +def runner_fixture( + request, + lxd: MockLxdClient, + jinja: MagicMock, + tmp_path: Path, + ssh_debug_infos: list[SSHDebugInfo], +): client = RunnerManagerClients( MagicMock(), jinja, @@ -117,6 +131,7 @@ def runner_fixture(request, lxd: MockLxdClient, jinja: MagicMock, tmp_path: Path lxd_storage_path=pool_path, dockerhub_mirror=None, issue_metrics=False, + ssh_debug_infos=ssh_debug_infos, ) status = RunnerStatus() return Runner( @@ -416,3 +431,32 @@ def test_remove_with_delete_error( with pytest.raises(RunnerRemoveError): runner.remove("test_token") + + +def test_random_ssh_connection_choice( + runner: Runner, + vm_resources: VirtualMachineResources, + token: str, + binary_path: Path, +): + """ + arrange: given a mock runner with random batch of ssh debug infos. + act: when runner.configure_runner is called. + assert: selected ssh_debug_info is random. + """ + runner.create( + config=CreateRunnerConfig( + image="test_image", + resources=vm_resources, + binary_path=binary_path, + registration_token=token, + ) + ) + runner._configure_runner() + first_call_args = runner._clients.jinja.get_template("env.j2").render.call_args.kwargs + runner._configure_runner() + second_call_args = runner._clients.jinja.get_template("env.j2").render.call_args.kwargs + + assert ( + first_call_args["ssh_debug_info"] != second_call_args["ssh_debug_info"] + ), "Same ssh debug info found, this may have occurred with a very low priority." diff --git a/tests/unit/test_runner_manager.py b/tests/unit/test_runner_manager.py index 4b7d3a037..2e778c939 100644 --- a/tests/unit/test_runner_manager.py +++ b/tests/unit/test_runner_manager.py @@ -35,7 +35,7 @@ def charm_state_fixture(): mock = MagicMock(spec=State) mock.is_metrics_logging_available = False mock.arch = ARCH.X64 - mock.ssh_debug_info = None + mock.ssh_debug_infos = None return mock diff --git a/tox.ini b/tox.ini index 9dfb8b1ce..a74a585ba 100644 --- a/tox.ini +++ b/tox.ini @@ -77,6 +77,7 @@ deps = requests-mock coverage[toml] -r{toxinidir}/requirements.txt + -r{[vars]tst_path}unit/requirements.txt commands = coverage run --source={[vars]src_path} \ -m pytest --ignore={[vars]tst_path}integration -v --tb native -s {posargs}