feat(agent): use typed_config in TF Agent Host Charm (#927)

* feat(agent): add typed config to agent host charm

* tests: update unit tests with typed config

* remove unecessary try except

* add additional unit test for invalid base64

* add reviewer suggestions

Author: rene-oromtz

agent/charms/testflinger-agent-host-charm/requirements.txt

@@ -7,4 +7,5 @@ Jinja2==3.1.6
 maturin==1.8.6
 charmlibs-passwd==1.0.0
 charmlibs-apt==1.0.0
-requests>=2.32.4
+requests>=2.32.4
+pydantic>=2.12,<3

agent/charms/testflinger-agent-host-charm/src/charm.py

@@ -18,6 +18,7 @@
 from charmlibs import apt, passwd
 from charms.grafana_agent.v0.cos_agent import COSAgentProvider
 from common import copy_ssh_keys, run_with_logged_errors, update_charm_scripts
+from config import TestflingerAgentConfig
 from defaults import (
     AGENT_CONFIGS_PATH,
     LOCAL_TESTFLINGER_PATH,
@@ -36,6 +37,9 @@ class TestflingerAgentHostCharm(ops.charm.CharmBase):
 
     def __init__(self, *args):
         super().__init__(*args)
+        self.typed_config = self.load_config(
+            TestflingerAgentConfig, errors="blocked"
+        )
         self.framework.observe(self.on.install, self.on_install)
         self.framework.observe(self.on.start, self.on_start)
         self.framework.observe(self.on.config_changed, self.on_config_changed)
@@ -106,10 +110,10 @@ def update_config_files(self):
         Clone the config files from the repo and swap it in for whatever is
         in AGENT_CONFIGS_PATH.
         """
-        config_repo = self.config.get("config-repo")
-        config_dir = self.config.get("config-dir")
-        config_branch = self.config.get("config-branch")
-        if not config_repo or not config_dir:
+        if (
+            not self.typed_config.config_repo
+            or not self.typed_config.config_dir
+        ):
             logger.error("config-repo and config-dir must be set")
             raise ValueError("config-repo and config-dir must be set")
         tmp_repo_path = Path("/srv/tmp-agent-configs")
@@ -118,8 +122,8 @@ def update_config_files(self):
             shutil.rmtree(tmp_repo_path, ignore_errors=True)
         try:
             Repo.clone_from(
-                url=config_repo,
-                branch=config_branch,
+                url=self.typed_config.config_repo,
+                branch=self.typed_config.config_branch,
                 to_path=tmp_repo_path,
                 depth=1,
             )
@@ -140,7 +144,7 @@ def write_supervisor_service_files(self):
         contains a directory for each agent that needs to run from this host.
         The agent directory name will be used as the service name.
         """
-        config_dirs = Path(AGENT_CONFIGS_PATH) / self.config.get("config-dir")
+        config_dirs = Path(AGENT_CONFIGS_PATH) / self.typed_config.config_dir
 
         if not config_dirs.is_dir():
             logger.error("config-dir must point to a directory")
@@ -219,7 +223,7 @@ def setup_docker(self):
     def update_tf_cmd_scripts(self):
         """Update tf-cmd-scripts."""
         self.unit.status = ops.MaintenanceStatus("Installing tf-cmd-scripts")
-        update_charm_scripts(self.config)
+        update_charm_scripts(self.typed_config)
 
     def on_upgrade_charm(self, _):
         """Upgrade hook."""
@@ -237,7 +241,8 @@ def on_start(self, _):
 
     def _on_secret_changed(self, event: ops.SecretChangedEvent):
         """Handle secret changed event."""
-        if event.secret.id == self.config.get("credentials-secret"):
+        secret = self.typed_config.credentials_secret
+        if secret and event.secret.id == secret.id:
             logger.info("Credentials secret changed, re-authenticating")
             self._update_refresh_token()
 
@@ -260,18 +265,16 @@ def _block(self, message: str) -> bool:
     def _valid_secret(self) -> dict | None:
         """Check the validity of the credentials-secret.
 
-        This evaluates if the secret URI is known and and if so that the secret
+        This evaluates if the secret is known and if so that the secret
         contains the necessary fields.
 
         If secret is valid, return the secret content, otherwise return None.
         """
-        secret_uri = self.config.get("credentials-secret")
-        if not secret_uri:
+        if not (secret := self.typed_config.credentials_secret):
             logger.error("credentials-secret config not set")
             return
 
         try:
-            secret = self.model.get_secret(id=secret_uri)
             content = secret.get_content(refresh=True)
             if "client-id" not in content or "secret-key" not in content:
                 logger.error("Secret missing required fields")
@@ -302,13 +305,9 @@ def _authenticate_with_server(self) -> bool:
         if not (content := self._valid_secret()):
             return self._block("Invalid credentials secret")
 
-        server = self.config.get("testflinger-server")
-        if not server or not server.startswith(("http://", "https://")):
-            return self._block("Testflinger server config not set or invalid")
-
         logger.info("Authenticating with Testflinger server")
         if not testflinger_client.authenticate(
-            server=server,
+            server=self.typed_config.testflinger_server,
             client_id=content["client-id"],
             secret_key=content["secret-key"],
         ):
@@ -325,7 +324,7 @@ def on_config_changed(self, _):
         except ValueError:
             self._block("config-repo and config-dir must be set")
             return
-        copy_ssh_keys(self.config)
+        copy_ssh_keys(self.typed_config)
         self.update_tf_cmd_scripts()
         self.write_supervisor_service_files()
         supervisord.supervisor_update()

agent/charms/testflinger-agent-host-charm/src/common.py

@@ -4,9 +4,9 @@
 import logging
 import os
 import subprocess
-from base64 import b64decode
 from pathlib import Path
 
+from config import TestflingerAgentConfig
 from defaults import (
     AGENT_CONFIGS_PATH,
     VIRTUAL_ENV_PATH,
@@ -34,24 +34,16 @@ def run_with_logged_errors(cmd: list) -> int:
     return proc.returncode
 
 
-def copy_ssh_keys(config: dict) -> None:
-    try:
-        ssh_config = config.get("ssh-config", "")
-        config_file = Path(SSH_CONFIG)
-        write_file(config_file, ssh_config, chmod=0o640)
+def copy_ssh_keys(config: TestflingerAgentConfig) -> None:
+    """Copy SSH keys and config from charm config to ~/.ssh/.
 
-        priv_key = config.get("ssh-private-key", "")
-        priv_key_file = Path(SSH_PRIVATE_KEY)
-        write_file(priv_key_file, b64decode(priv_key).decode(), chmod=0o600)
+    :param config: The charm configuration containing the SSH keys and config.
+    """
+    write_file(Path(SSH_CONFIG), config.ssh_config, chmod=0o640)
 
-        pub_key = config.get("ssh-public-key", "")
-        pub_key_file = Path(SSH_PUBLIC_KEY)
-        write_file(pub_key_file, b64decode(pub_key).decode())
-    except (TypeError, UnicodeDecodeError):
-        logger.error(
-            "Failed to decode ssh keys - ensure they are base64 encoded"
-        )
-        raise
+    write_file(Path(SSH_PRIVATE_KEY), config.ssh_private_key, chmod=0o600)
+
+    write_file(Path(SSH_PUBLIC_KEY), config.ssh_public_key)
 
 
 def write_file(location: Path, contents: str, chmod: int = 0o644) -> None:
@@ -69,15 +61,19 @@ def write_file(location: Path, contents: str, chmod: int = 0o644) -> None:
     os.chmod(location, chmod)
 
 
-def update_charm_scripts(config: dict) -> None:
+def update_charm_scripts(config: TestflingerAgentConfig) -> None:
+    """Update charm scripts with rendered templates for tf scripts.
+
+    :param config: The charm configuration containing config directory.
+    """
     tf_cmd_dir = Path("src/tf-cmd-scripts/")
     usr_local_bin = Path("/usr/local/bin")
 
     for tf_cmd_file in tf_cmd_dir.iterdir():
         template = Template(tf_cmd_file.read_text())
         rendered = template.render(
             agent_configs_path=AGENT_CONFIGS_PATH,
-            config_dir=config.get("config-dir"),
+            config_dir=config.config_dir,
             virtual_env_path=VIRTUAL_ENV_PATH,
         )
         agent_file = usr_local_bin / tf_cmd_file.name

agent/charms/testflinger-agent-host-charm/src/config.py

@@ -0,0 +1,37 @@
+# Copyright 2026 Canonical Ltd.
+# See LICENSE file for licensing details.
+"""Charm configuration."""
+
+import ops
+import pydantic
+
+SSH_CONFIG_FILE = """\
+StrictHostKeyChecking no
+UserKnownHostsFile /dev/null
+LogLevel QUIET
+ConnectTimeout 30
+"""
+
+
+class TestflingerAgentConfig(pydantic.BaseModel):
+    """Testflinger Agent Host Charm configuration."""
+
+    model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+
+    config_repo: str = ""
+    config_branch: str = "main"
+    config_dir: str = ""
+    ssh_private_key: pydantic.Base64Str = ""
+    ssh_public_key: pydantic.Base64Str = ""
+    ssh_config: str = SSH_CONFIG_FILE
+    testflinger_server: str = "https://testflinger.canonical.com"
+    credentials_secret: ops.Secret | None = None
+
+    @pydantic.field_validator("testflinger_server")
+    @classmethod
+    def validate_server(cls, value):
+        if not value.startswith(("http://", "https://")):
+            raise ValueError(
+                "testflinger_server must include protocol (http:// or https://)"
+            )
+        return value

agent/charms/testflinger-agent-host-charm/tests/unit/test_charm.py

@@ -49,21 +49,6 @@ def test_update_tf_cmd_scripts_on_config_changed(
     mock_update_scripts.assert_called_once()
 
 
-def test_blocked_on_invalid_server(ctx, state_in, secret):
-    """Test blocked status when server is not valid."""
-    state = state_in(
-        config={
-            "testflinger-server": "localhost:5000",
-            "credentials-secret": secret.id,
-        },
-        secrets=[secret],
-    )
-    state_out = ctx.run(ctx.on.start(), state=state)
-    assert state_out.unit_status == testing.BlockedStatus(
-        "Testflinger server config not set or invalid"
-    )
-
-
 def test_blocked_on_no_secret(ctx, state_in, caplog):
     """Test blocked status when credentials-secret config is not set."""
     state_out = ctx.run(ctx.on.start(), state=state_in())

agent/charms/testflinger-agent-host-charm/tests/unit/test_common.py

@@ -6,6 +6,7 @@
 from unittest.mock import MagicMock, patch
 
 import common
+from config import TestflingerAgentConfig
 
 SSH_PUBLIC_KEY = "/home/ubuntu/.ssh/id_rsa.pub"
 SSH_PRIVATE_KEY = "/home/ubuntu/.ssh/id_rsa"
@@ -16,13 +17,11 @@
 @patch("common.write_file")
 def test_copy_ssh_keys(mock_write_file, mock_chmod, mock_chown):
     """Test the copy_ssh_keys function."""
-    config = {
-        "ssh-config": "ssh_config_content",
-        "ssh-private-key": base64.b64encode(
-            b"ssh_private_key_content"
-        ).decode(),
-        "ssh-public-key": base64.b64encode(b"ssh_public_key_content").decode(),
-    }
+    config = TestflingerAgentConfig(
+        ssh_config="ssh_config_content",
+        ssh_private_key=base64.b64encode(b"ssh_private_key_content").decode(),
+        ssh_public_key=base64.b64encode(b"ssh_public_key_content").decode(),
+    )
 
     common.copy_ssh_keys(config)
 
@@ -51,7 +50,7 @@ def test_update_charm_scripts(
     mock_script.read_text.return_value = "config_dir={{ config_dir }}"
     mock_iterdir.return_value = [mock_script]
 
-    config = {"config-dir": "my-agent-configs"}
+    config = TestflingerAgentConfig(config_dir="my-agent-configs")
 
     common.update_charm_scripts(config)
 

agent/charms/testflinger-agent-host-charm/tests/unit/test_config.py

@@ -0,0 +1,48 @@
+# Copyright 2026 Canonical
+# See LICENSE file for licensing details.
+"""Unit tests for config.py."""
+
+import base64
+
+import pytest
+from config import TestflingerAgentConfig
+
+
+def test_base64_str_fields():
+    """Test that Base64Str fields decode correctly."""
+    private_key = base64.b64encode(b"ssh_private_key_content").decode()
+    public_key = base64.b64encode(b"ssh_public_key_content").decode()
+    config = TestflingerAgentConfig(
+        ssh_private_key=private_key,
+        ssh_public_key=public_key,
+    )
+    assert config.ssh_private_key == "ssh_private_key_content"
+    assert config.ssh_public_key == "ssh_public_key_content"
+
+
+def test_invalid_base64_str_fields():
+    """Test that invalid Base64Str fields raise errors."""
+    with pytest.raises(ValueError):
+        TestflingerAgentConfig(ssh_private_key="undecodable_string")
+
+    with pytest.raises(ValueError):
+        TestflingerAgentConfig(ssh_public_key="undecodable_string")
+
+
+def test_valid_testflinger_server():
+    """Test that valid server URLs are accepted."""
+    config = TestflingerAgentConfig(
+        testflinger_server="https://testflinger.local"
+    )
+    assert config.testflinger_server == "https://testflinger.local"
+
+    config = TestflingerAgentConfig(
+        testflinger_server="http://testflinger.local"
+    )
+    assert config.testflinger_server == "http://testflinger.local"
+
+
+def test_invalid_testflinger_server():
+    """Test that invalid server URLs are rejected."""
+    with pytest.raises(ValueError):
+        TestflingerAgentConfig(testflinger_server="localhost:5000")