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

@@ -183,7 +183,7 @@ SSH connection information for debug workflow.
 ### <kbd>classmethod</kbd> `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.
  - <b>`proxy_config`</b>:  Proxy-related configuration. 
  - <b>`charm_config`</b>:  Configuration of the juju charm. 
  - <b>`arch`</b>:  The underlying compute architecture, i.e. x86_64, amd64, arm64/aarch64. 
- - <b>`ssh_debug_info`</b>:  The SSH debug connection configuration information. 
+ - <b>`ssh_debug_infos`</b>:  SSH debug connections configuration information. 
 
 
 
 
 ---
 
-<a href="../src/charm_state.py#L248"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
+<a href="../src/charm_state.py#L254"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
 
 ### <kbd>classmethod</kbd> `from_charm`
 

src-docs/runner.py.md

@@ -39,7 +39,7 @@ The configuration values for creating a single runner instance.
 ## <kbd>class</kbd> `Runner`
 Single instance of GitHub self-hosted runner. 
 
-<a href="../src/runner.py#L101"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
+<a href="../src/runner.py#L102"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
 
 ### <kbd>function</kbd> `__init__`
 
@@ -67,7 +67,7 @@ Construct the runner instance.
 
 ---
 
-<a href="../src/runner.py#L131"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
+<a href="../src/runner.py#L132"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
 
 ### <kbd>function</kbd> `create`
 
@@ -91,7 +91,7 @@ Create the runner instance on LXD and register it on GitHub.
 
 ---
 
-<a href="../src/runner.py#L167"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
+<a href="../src/runner.py#L168"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
 
 ### <kbd>function</kbd> `remove`
 

src-docs/runner_type.py.md

@@ -83,7 +83,7 @@ Configuration for runner.
  - <b>`path`</b>:  GitHub repository path in the format '<owner>/<repo>', or the GitHub organization  name. 
  - <b>`proxies`</b>:  HTTP(S) proxy settings. 
  - <b>`dockerhub_mirror`</b>:  URL of dockerhub mirror to use. 
- - <b>`ssh_debug_info`</b>:  The SSH debug server connection metadata. 
+ - <b>`ssh_debug_infos`</b>:  The SSH debug server connections metadata. 
 
 
 

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")
 

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)])

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,

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

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:")

tests/unit/requirements.txt

@@ -0,0 +1 @@
+factory-boy>=3,<4

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():