Skip to content

Commit 14ccdf5

Browse files
dragomirpdragomirp
committed
Factor out helper and unit tests
1 parent cc6c8fa commit 14ccdf5

3 files changed

Lines changed: 103 additions & 59 deletions

File tree

‎src/charm.py‎

Lines changed: 9 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -4,63 +4,6 @@
44

55
"""Charmed Machine Operator for the PostgreSQL database."""
66

7-
8-
def _remove_stale_otel_sdk_packages():
9-
"""Hack to remove stale opentelemetry sdk packages from the charm's python venv.
10-
11-
See https://github.com/canonical/grafana-agent-operator/issues/146 and
12-
https://bugs.launchpad.net/juju/+bug/2058335 for more context. This patch can be removed after
13-
this juju issue is resolved and sufficient time has passed to expect most users of this library
14-
have migrated to the patched version of juju. When this patch is removed, un-ignore rule E402 for this file in the pyproject.toml (see setting
15-
[tool.ruff.lint.per-file-ignores] in pyproject.toml).
16-
17-
This only has an effect if executed on an upgrade-charm event.
18-
"""
19-
# all imports are local to keep this function standalone, side-effect-free, and easy to revert later
20-
import os
21-
22-
if os.getenv("JUJU_DISPATCH_PATH") != "hooks/upgrade-charm":
23-
return
24-
25-
import logging
26-
import shutil
27-
from collections import defaultdict
28-
29-
from importlib_metadata import distributions
30-
31-
otel_logger = logging.getLogger("charm_tracing_otel_patcher")
32-
otel_logger.debug("Applying _remove_stale_otel_sdk_packages patch on charm upgrade")
33-
# group by name all distributions starting with "opentelemetry_"
34-
otel_distributions = defaultdict(list)
35-
for distribution in distributions():
36-
name = distribution._normalized_name
37-
if name.startswith("opentelemetry_"):
38-
otel_distributions[name].append(distribution)
39-
40-
otel_logger.debug(f"Found {len(otel_distributions)} opentelemetry distributions")
41-
42-
# If we have multiple distributions with the same name, remove any that have 0 associated files
43-
for name, distributions_ in otel_distributions.items():
44-
if len(distributions_) <= 1:
45-
continue
46-
47-
otel_logger.debug(f"Package {name} has multiple ({len(distributions_)}) distributions.")
48-
for distribution in distributions_:
49-
if not distribution.files: # Not None or empty list
50-
path = distribution._path
51-
otel_logger.info(f"Removing empty distribution of {name} at {path}.")
52-
shutil.rmtree(path)
53-
54-
otel_logger.debug("Successfully applied _remove_stale_otel_sdk_packages patch. ")
55-
56-
57-
# apply hacky patch to remove stale opentelemetry sdk packages on upgrade-charm.
58-
# it could be trouble if someone ever decides to implement their own tracer parallel to
59-
# ours and before the charm has inited. We assume they won't.
60-
# !!IMPORTANT!! keep all otlp imports UNDER this call.
61-
_remove_stale_otel_sdk_packages()
62-
63-
# ruff: disable[E402]
647
import json
658
import logging
669
import os
@@ -76,6 +19,15 @@ def _remove_stale_otel_sdk_packages():
7619
from typing import Literal, get_args
7720
from urllib.parse import urlparse
7821

22+
from utils import _remove_stale_otel_sdk_packages
23+
24+
# apply hacky patch to remove stale opentelemetry sdk packages on upgrade-charm.
25+
# it could be trouble if someone ever decides to implement their own tracer parallel to
26+
# ours and before the charm has inited. We assume they won't.
27+
# !!IMPORTANT!! keep all otlp imports UNDER this call.
28+
_remove_stale_otel_sdk_packages()
29+
30+
# ruff: disable[E402]
7931
import psycopg2
8032
from charmlibs import snap
8133
from charms.data_platform_libs.v0.data_interfaces import DataPeerData, DataPeerUnitData

‎src/utils.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,15 @@
33

44
"""A collection of utility functions that are used in the charm."""
55

6+
import logging
7+
import os
68
import platform
79
import secrets
10+
import shutil
811
import string
12+
from collections import defaultdict
13+
14+
from importlib_metadata import distributions
915

1016
from constants import (
1117
POSTGRESQL_SNAP_NAME,
@@ -47,3 +53,45 @@ def label2name(label: str) -> str:
4753
The converted name.
4854
"""
4955
return label.rsplit("-", 1)[0] + "/" + label.rsplit("-", 1)[1]
56+
57+
58+
def _remove_stale_otel_sdk_packages():
59+
"""Hack to remove stale opentelemetry sdk packages from the charm's python venv.
60+
61+
See https://github.com/canonical/grafana-agent-operator/issues/146 and
62+
https://bugs.launchpad.net/juju/+bug/2058335 for more context. This patch can be removed after
63+
this juju issue is resolved and sufficient time has passed to expect most users of this library
64+
have migrated to the patched version of juju. When this patch is removed, un-ignore rule E402 for this file in the pyproject.toml (see setting
65+
[tool.ruff.lint.per-file-ignores] in pyproject.toml).
66+
67+
This only has an effect if executed on an upgrade-charm event.
68+
"""
69+
# all imports are local to keep this function standalone, side-effect-free, and easy to revert later
70+
71+
if os.getenv("JUJU_DISPATCH_PATH") != "hooks/upgrade-charm":
72+
return
73+
74+
otel_logger = logging.getLogger("charm_tracing_otel_patcher")
75+
otel_logger.debug("Applying _remove_stale_otel_sdk_packages patch on charm upgrade")
76+
# group by name all distributions starting with "opentelemetry_"
77+
otel_distributions = defaultdict(list)
78+
for distribution in distributions():
79+
name = distribution._normalized_name
80+
if name.startswith("opentelemetry_"):
81+
otel_distributions[name].append(distribution)
82+
83+
otel_logger.debug(f"Found {len(otel_distributions)} opentelemetry distributions")
84+
85+
# If we have multiple distributions with the same name, remove any that have 0 associated files
86+
for name, distributions_ in otel_distributions.items():
87+
if len(distributions_) <= 1:
88+
continue
89+
90+
otel_logger.debug(f"Package {name} has multiple ({len(distributions_)}) distributions.")
91+
for distribution in distributions_:
92+
if not distribution.files: # Not None or empty list
93+
path = distribution._path
94+
otel_logger.info(f"Removing empty distribution of {name} at {path}.")
95+
shutil.rmtree(path)
96+
97+
otel_logger.debug("Successfully applied _remove_stale_otel_sdk_packages patch.")

‎tests/unit/test_utils.py‎

Lines changed: 46 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,10 @@
22
# See LICENSE file for licensing details.
33

44
import re
5-
from unittest.mock import patch
5+
from unittest.mock import Mock, patch, sentinel
66

77
from constants import POSTGRESQL_SNAP_NAME
8-
from utils import new_password, snap_refreshed
8+
from utils import _remove_stale_otel_sdk_packages, new_password, snap_refreshed
99

1010

1111
def test_new_password():
@@ -34,3 +34,47 @@ def test_snap_refreshed():
3434
):
3535
assert snap_refreshed("100") is False
3636
assert snap_refreshed("200") is False
37+
38+
39+
def test_remove_stale_otel_sdk_packages():
40+
with (
41+
patch("utils.os.getenv", return_value=None) as _getenv,
42+
patch("utils.shutil") as _shutil,
43+
patch("utils.distributions") as _distributions,
44+
):
45+
other_dist = Mock()
46+
other_dist._normalized_name = "test"
47+
otel_dist = Mock()
48+
otel_dist._normalized_name = "opentelemetry_test"
49+
stale_otel_dist = Mock()
50+
stale_otel_dist._normalized_name = "opentelemetry_test"
51+
stale_otel_dist.files = []
52+
stale_otel_dist._path = sentinel.path
53+
54+
# Not called if not upgrade hook
55+
_remove_stale_otel_sdk_packages()
56+
57+
_distributions.assert_not_called()
58+
_shutil.rmtree.assert_not_called()
59+
_distributions.reset_mock()
60+
_shutil.rmtree.reset_mock()
61+
62+
# Upgrade hook, nothing to remove
63+
_getenv.return_value = "hooks/upgrade-charm"
64+
_distributions.return_value = [other_dist, otel_dist]
65+
66+
_remove_stale_otel_sdk_packages()
67+
68+
_distributions.assert_called_once_with()
69+
_shutil.rmtree.assert_not_called()
70+
_distributions.reset_mock()
71+
_shutil.rmtree.reset_mock()
72+
73+
# Upgrade hook, duplicate otel packages
74+
_distributions.return_value = [other_dist, otel_dist, stale_otel_dist]
75+
_remove_stale_otel_sdk_packages()
76+
77+
_distributions.assert_called_once_with()
78+
_shutil.rmtree.assert_called_once_with(sentinel.path)
79+
_distributions.reset_mock()
80+
_shutil.rmtree.reset_mock()

0 commit comments

Comments
 (0)