|  | 
| 9 | 9 | import psycopg2 | 
| 10 | 10 | import requests | 
| 11 | 11 | import yaml | 
|  | 12 | +from lightkube import codecs | 
|  | 13 | +from lightkube.core.client import Client | 
|  | 14 | +from lightkube.core.exceptions import ApiError | 
|  | 15 | +from lightkube.generic_resource import GenericNamespacedResource | 
|  | 16 | +from lightkube.resources.core_v1 import Endpoints, Service | 
| 12 | 17 | from pytest_operator.plugin import OpsTest | 
| 13 | 18 | from tenacity import retry, retry_if_result, stop_after_attempt, wait_exponential | 
| 14 | 19 | 
 | 
| @@ -222,6 +227,110 @@ def get_application_units(ops_test: OpsTest, application_name: str) -> List[str] | 
| 222 | 227 |     ] | 
| 223 | 228 | 
 | 
| 224 | 229 | 
 | 
|  | 230 | +def get_charm_resources(namespace: str, application: str) -> List[GenericNamespacedResource]: | 
|  | 231 | +    """Return the list of k8s resources from resources.yaml file. | 
|  | 232 | +
 | 
|  | 233 | +    Args: | 
|  | 234 | +        namespace: namespace related to the model where | 
|  | 235 | +            the charm was deployed. | 
|  | 236 | +        application: application name. | 
|  | 237 | +
 | 
|  | 238 | +    Returns: | 
|  | 239 | +        list of existing charm/Patroni specific k8s resources. | 
|  | 240 | +    """ | 
|  | 241 | +    # Define the context needed for the k8s resources lists load. | 
|  | 242 | +    context = {"namespace": namespace, "app_name": application} | 
|  | 243 | + | 
|  | 244 | +    # Load the list of the resources from resources.yaml. | 
|  | 245 | +    with open("src/resources.yaml") as f: | 
|  | 246 | +        return codecs.load_all_yaml(f, context=context) | 
|  | 247 | + | 
|  | 248 | + | 
|  | 249 | +def get_existing_k8s_resources(namespace: str, application: str) -> set: | 
|  | 250 | +    """Return the list of k8s resources that were created by the charm and Patroni. | 
|  | 251 | +
 | 
|  | 252 | +    Args: | 
|  | 253 | +        namespace: namespace related to the model where | 
|  | 254 | +            the charm was deployed. | 
|  | 255 | +        application: application name. | 
|  | 256 | +
 | 
|  | 257 | +    Returns: | 
|  | 258 | +        list of existing charm/Patroni specific k8s resources. | 
|  | 259 | +    """ | 
|  | 260 | +    # Create a k8s API client instance. | 
|  | 261 | +    client = Client(namespace=namespace) | 
|  | 262 | + | 
|  | 263 | +    # Retrieve the k8s resources the charm should create. | 
|  | 264 | +    charm_resources = get_charm_resources(namespace, application) | 
|  | 265 | + | 
|  | 266 | +    # Add only the resources that currently exist. | 
|  | 267 | +    resources = set( | 
|  | 268 | +        map( | 
|  | 269 | +            # Build an identifier for each resource (using its type and name). | 
|  | 270 | +            lambda x: f"{type(x).__name__}/{x.metadata.name}", | 
|  | 271 | +            filter( | 
|  | 272 | +                lambda x: (resource_exists(client, x)), | 
|  | 273 | +                charm_resources, | 
|  | 274 | +            ), | 
|  | 275 | +        ) | 
|  | 276 | +    ) | 
|  | 277 | + | 
|  | 278 | +    # Include the resources created by the charm and Patroni. | 
|  | 279 | +    for kind in [Endpoints, Service]: | 
|  | 280 | +        extra_resources = client.list( | 
|  | 281 | +            kind, | 
|  | 282 | +            namespace=namespace, | 
|  | 283 | +            labels={"app.juju.is/created-by": application}, | 
|  | 284 | +        ) | 
|  | 285 | +        resources.update( | 
|  | 286 | +            set( | 
|  | 287 | +                map( | 
|  | 288 | +                    # Build an identifier for each resource (using its type and name). | 
|  | 289 | +                    lambda x: f"{kind.__name__}/{x.metadata.name}", | 
|  | 290 | +                    extra_resources, | 
|  | 291 | +                ) | 
|  | 292 | +            ) | 
|  | 293 | +        ) | 
|  | 294 | + | 
|  | 295 | +    return resources | 
|  | 296 | + | 
|  | 297 | + | 
|  | 298 | +def get_expected_k8s_resources(namespace: str, application: str) -> set: | 
|  | 299 | +    """Return the list of expected k8s resources when the charm is deployed. | 
|  | 300 | +
 | 
|  | 301 | +    Args: | 
|  | 302 | +        namespace: namespace related to the model where | 
|  | 303 | +            the charm was deployed. | 
|  | 304 | +        application: application name. | 
|  | 305 | +
 | 
|  | 306 | +    Returns: | 
|  | 307 | +        list of existing charm/Patroni specific k8s resources. | 
|  | 308 | +    """ | 
|  | 309 | +    # Retrieve the k8s resources created by the charm. | 
|  | 310 | +    charm_resources = get_charm_resources(namespace, application) | 
|  | 311 | + | 
|  | 312 | +    # Build an identifier for each resource (using its type and name). | 
|  | 313 | +    resources = set( | 
|  | 314 | +        map( | 
|  | 315 | +            lambda x: f"{type(x).__name__}/{x.metadata.name}", | 
|  | 316 | +            charm_resources, | 
|  | 317 | +        ) | 
|  | 318 | +    ) | 
|  | 319 | + | 
|  | 320 | +    # Include the resources created by the charm and Patroni. | 
|  | 321 | +    resources.update( | 
|  | 322 | +        [ | 
|  | 323 | +            f"Endpoints/patroni-{application}-config", | 
|  | 324 | +            f"Endpoints/patroni-{application}", | 
|  | 325 | +            f"Endpoints/{application}-primary", | 
|  | 326 | +            f"Endpoints/{application}-replicas", | 
|  | 327 | +            f"Service/patroni-{application}-config", | 
|  | 328 | +        ] | 
|  | 329 | +    ) | 
|  | 330 | + | 
|  | 331 | +    return resources | 
|  | 332 | + | 
|  | 333 | + | 
| 225 | 334 | async def get_password(ops_test: OpsTest, username: str = "operator"): | 
| 226 | 335 |     """Retrieve a user password using the action.""" | 
| 227 | 336 |     unit = ops_test.model.units.get(f"{DATABASE_APP_NAME}/0") | 
| @@ -272,6 +381,23 @@ async def restart_patroni(ops_test: OpsTest, unit_name: str) -> None: | 
| 272 | 381 |     requests.post(f"http://{unit_ip}:8008/restart") | 
| 273 | 382 | 
 | 
| 274 | 383 | 
 | 
|  | 384 | +def resource_exists(client: Client, resource: GenericNamespacedResource) -> bool: | 
|  | 385 | +    """Check whether a specific resource exists. | 
|  | 386 | +
 | 
|  | 387 | +    Args: | 
|  | 388 | +        client: k8s API client instance. | 
|  | 389 | +        resource: k8s resource. | 
|  | 390 | +
 | 
|  | 391 | +    Returns: | 
|  | 392 | +        whether the resource exists. | 
|  | 393 | +    """ | 
|  | 394 | +    try: | 
|  | 395 | +        client.get(type(resource), name=resource.metadata.name) | 
|  | 396 | +        return True | 
|  | 397 | +    except ApiError: | 
|  | 398 | +        return False | 
|  | 399 | + | 
|  | 400 | + | 
| 275 | 401 | async def scale_application(ops_test: OpsTest, application_name: str, scale: int) -> None: | 
| 276 | 402 |     """Scale a given application to a specific unit count. | 
| 277 | 403 | 
 | 
|  | 
0 commit comments