From 01cf131568a629355f2e0d851a6deaa34f7d47c1 Mon Sep 17 00:00:00 2001
From: Adam Dyess <adam.dyess@canonical.com>
Date: Mon, 16 Dec 2024 12:40:14 -0600
Subject: [PATCH] Render CharmURL to yaml

---
 tests/integration/helpers.py | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/tests/integration/helpers.py b/tests/integration/helpers.py
index 03a22054..eacf326f 100644
--- a/tests/integration/helpers.py
+++ b/tests/integration/helpers.py
@@ -14,7 +14,7 @@
 from functools import cache, cached_property
 from itertools import chain
 from pathlib import Path
-from typing import Dict, List, Mapping, Optional, Tuple
+from typing import Dict, List, Mapping, Optional, Set, Tuple
 
 import yaml
 from juju import unit
@@ -267,13 +267,19 @@ def craft(cls, name: str, series: str, arch: str) -> "CharmUrl":
             name = m.group("charm")
         return cls(name, series, arch)
 
-    def __str__(self) -> str:
-        """Return the charm URL.
+    @staticmethod
+    def representer(dumper: yaml.Dumper, data: "CharmUrl") -> yaml.ScalarNode:
+        """Yaml representer for the CharmUrl object.
+
+        Args:
+            dumper: yaml dumper
+            data: CharmUrl object
 
         Returns:
-            string: charm URL
+            yaml.ScalarNode: yaml node
         """
-        return f"ch:{self.arch}/{self.series}/{self.name}"
+        repr = f"ch:{data.arch}/{data.series}/{data.name}"
+        return dumper.represent_scalar("tag:yaml.org,2002:str", repr)
 
 
 @dataclass
@@ -423,7 +429,7 @@ def content(self) -> Mapping:
             loaded = yaml.safe_load(self.path.read_bytes())
             series = loaded.get("series", "focal")
             for app in loaded["applications"].values():
-                app["charm"] = CharmUrl(app["charm"], series=series, arch=self.arch)
+                app["charm"] = CharmUrl.craft(app["charm"], series=series, arch=self.arch)
             self._content = loaded
         return self._content
 
@@ -544,7 +550,7 @@ def render(self, tmp_path: Path) -> Path:
         """
         target = tmp_path / "bundles" / self.path.name
         target.parent.mkdir(exist_ok=True, parents=True)
-        yaml.safe_dump(self.content, target.open("w"))
+        yaml.dump(self.content, target.open("w"))
         return target
 
 
@@ -560,11 +566,11 @@ async def cloud_arch(ops_test: OpsTest) -> str:
     assert ops_test.model, "Model must be present"
     controller = await ops_test.model.get_controller()
     controller_model = await controller.get_model("controller")
-    arch = set(
+    arch: Set[str] = {
         machine.safe_data["hardware-characteristics"]["arch"]
         for machine in controller_model.machines.values()
-    )
-    return arch.pop()
+    }
+    return arch.pop().strip()
 
 
 async def cloud_type(ops_test: OpsTest) -> Tuple[str, bool]:
@@ -586,3 +592,6 @@ async def cloud_type(ops_test: OpsTest) -> Tuple[str, bool]:
     if _type == "lxd":
         vms = not ops_test.request.config.getoption("--lxd-containers")
     return _type, vms
+
+
+yaml.add_representer(CharmUrl, CharmUrl.representer)