[DPE-1697] Merge test apps (#113)

* Switch to merged app

* Use pgb k8s merged app

* Fixes

* Missing arg

* Missing arg

* Update to PG VM test app

* Switch to charmhub test app

* Bump libs

lib/charms/grafana_agent/v0/cos_agent.py

@@ -236,9 +236,9 @@ class _MetricsEndpointDict(TypedDict):
 
 LIBID = "dc15fa84cef84ce58155fb84f6c6213a"
 LIBAPI = 0
-LIBPATCH = 5
+LIBPATCH = 6
 
-PYDEPS = ["cosl", "pydantic<2"]
+PYDEPS = ["cosl", "pydantic < 2"]
 
 DEFAULT_RELATION_NAME = "cos-agent"
 DEFAULT_PEER_RELATION_NAME = "peers"
@@ -285,6 +285,7 @@ class CosAgentProviderUnitData(pydantic.BaseModel):
     metrics_alert_rules: dict
     log_alert_rules: dict
     dashboards: List[GrafanaDashboard]
+    subordinate: Optional[bool]
 
     # The following entries may vary across units of the same principal app.
     # this data does not need to be forwarded to the gagent leader
@@ -401,6 +402,7 @@ def _on_refresh(self, event):
                         dashboards=self._dashboards,
                         metrics_scrape_jobs=self._scrape_jobs,
                         log_slots=self._log_slots,
+                        subordinate=self._charm.meta.subordinate,
                     )
                     relation.data[self._charm.unit][data.KEY] = data.json()
                 except (
@@ -493,6 +495,12 @@ class COSAgentRequirerEvents(ObjectEvents):
     validation_error = EventSource(COSAgentValidationError)
 
 
+class MultiplePrincipalsError(Exception):
+    """Custom exception for when there are multiple principal applications."""
+
+    pass
+
+
 class COSAgentRequirer(Object):
     """Integration endpoint wrapper for the Requirer side of the cos_agent interface."""
 
@@ -589,7 +597,9 @@ def _on_relation_data_changed(self, event: RelationChangedEvent):
             log_alert_rules=provider_data.log_alert_rules,
             dashboards=provider_data.dashboards,
         )
-        self.peer_relation.data[self._charm.unit][data.KEY] = data.json()
+        self.peer_relation.data[self._charm.unit][
+            f"{CosAgentPeersUnitData.KEY}-{event.unit.name}"
+        ] = data.json()
 
         # We can't easily tell if the data that was changed is limited to only the data
         # that goes into peer relation (in which case, if this is not a leader unit, we wouldn't
@@ -627,35 +637,40 @@ def _principal_unit(self) -> Optional[Unit]:
 
     @property
     def _principal_relations(self):
-        # Technically it's a list, but for subordinates there can only be one.
-        return self._charm.model.relations[self._relation_name]
+        relations = []
+        for relation in self._charm.model.relations[self._relation_name]:
+            if not json.loads(relation.data[next(iter(relation.units))]["config"]).get(
+                ["subordinate"], False
+            ):
+                relations.append(relation)
+        if len(relations) > 1:
+            logger.error(
+                "Multiple applications claiming to be principal. Update the cos-agent library in the client application charms."
+            )
+            raise MultiplePrincipalsError("Multiple principal applications.")
+        return relations
 
     @property
-    def _principal_unit_data(self) -> Optional[CosAgentProviderUnitData]:
-        """Return the principal unit's data.
+    def _remote_data(self) -> List[CosAgentProviderUnitData]:
+        """Return a list of remote data from each of the related units.
 
         Assumes that the relation is of type subordinate.
         Relies on the fact that, for subordinate relations, the only remote unit visible to
         *this unit* is the principal unit that this unit is attached to.
         """
-        if not (relations := self._principal_relations):
-            return None
-
-        # Technically it's a list, but for subordinates there can only be one relation
-        principal_relation = next(iter(relations))
+        all_data = []
 
-        if not (units := principal_relation.units):
-            return None
-
-        # Technically it's a list, but for subordinates there can only be one
-        unit = next(iter(units))
-        if not (raw := principal_relation.data[unit].get(CosAgentProviderUnitData.KEY)):
-            return None
-
-        if not (provider_data := self._validated_provider_data(raw)):
-            return None
+        for relation in self._charm.model.relations[self._relation_name]:
+            if not relation.units:
+                continue
+            unit = next(iter(relation.units))
+            if not (raw := relation.data[unit].get(CosAgentProviderUnitData.KEY)):
+                continue
+            if not (provider_data := self._validated_provider_data(raw)):
+                continue
+            all_data.append(provider_data)
 
-        return provider_data
+        return all_data
 
     def _gather_peer_data(self) -> List[CosAgentPeersUnitData]:
         """Collect data from the peers.
@@ -673,18 +688,21 @@ def _gather_peer_data(self) -> List[CosAgentPeersUnitData]:
         app_names: Set[str] = set()
 
         for unit in chain((self._charm.unit,), relation.units):
-            if not relation.data.get(unit) or not (
-                raw := relation.data[unit].get(CosAgentPeersUnitData.KEY)
-            ):
-                logger.info(f"peer {unit} has not set its primary data yet; skipping for now...")
+            if not relation.data.get(unit):
                 continue
 
-            data = CosAgentPeersUnitData(**json.loads(raw))
-            app_name = data.app_name
-            # Have we already seen this principal app?
-            if app_name in app_names:
-                continue
-            peer_data.append(data)
+            for unit_name in relation.data.get(unit):  # pyright: ignore
+                if not unit_name.startswith(CosAgentPeersUnitData.KEY):
+                    continue
+                raw = relation.data[unit].get(unit_name)
+                if raw is None:
+                    continue
+                data = CosAgentPeersUnitData(**json.loads(raw))
+                # Have we already seen this principal app?
+                if (app_name := data.app_name) in app_names:
+                    continue
+                peer_data.append(data)
+                app_names.add(app_name)
 
         return peer_data
 
@@ -720,7 +738,7 @@ def metrics_alerts(self) -> Dict[str, Any]:
     def metrics_jobs(self) -> List[Dict]:
         """Parse the relation data contents and extract the metrics jobs."""
         scrape_jobs = []
-        if data := self._principal_unit_data:
+        for data in self._remote_data:
             for job in data.metrics_scrape_jobs:
                 # In #220, relation schema changed from a simplified dict to the standard
                 # `scrape_configs`.
@@ -740,7 +758,7 @@ def metrics_jobs(self) -> List[Dict]:
     def snap_log_endpoints(self) -> List[SnapEndpoint]:
         """Fetch logging endpoints exposed by related snaps."""
         plugs = []
-        if data := self._principal_unit_data:
+        for data in self._remote_data:
             targets = data.log_slots
             if targets:
                 for target in targets: