diff --git a/sarc/cli/acquire/storages.py b/sarc/cli/acquire/storages.py index 84df6cae..fefda5a9 100644 --- a/sarc/cli/acquire/storages.py +++ b/sarc/cli/acquire/storages.py @@ -1,3 +1,4 @@ +import logging from dataclasses import dataclass from simple_parsing import field @@ -6,6 +7,7 @@ from sarc.storage.diskusage import get_diskusage_collection from sarc.storage.drac import fetch_diskusage_report as fetch_dirac_diskusage from sarc.storage.mila import fetch_diskusage_report as fetch_mila_diskusage +from sarc.traces import using_trace methods = { "default": fetch_dirac_diskusage, @@ -22,18 +24,23 @@ def execute(self) -> int: cfg = config() for cluster_name in self.cluster_names: - print(f"Acquiring {cluster_name} storages report...") + with using_trace("AcquireStorages", "cluster") as span: + span.set_attribute("cluster_name", cluster_name) - cluster = cfg.clusters[cluster_name] + # pylint: disable=logging-fstring-interpolation + logging.info(f"Acquiring {cluster_name} storages report...") + span.add_event(f"Acquiring {cluster_name} storages report...") - fetch_diskusage = methods.get(cluster_name, methods["default"]) - du = fetch_diskusage(cluster) + cluster = cfg.clusters[cluster_name] - if not self.dry: - collection = get_diskusage_collection() - collection.add(du) - else: - print("Document:") - print(du.json(indent=2)) + fetch_diskusage = methods.get(cluster_name, methods["default"]) + du = fetch_diskusage(cluster) + + if not self.dry: + collection = get_diskusage_collection() + collection.add(du) + else: + print("Document:") + print(du.json(indent=2)) return 0 diff --git a/tests/functional/diskusage/test_mila_diskusage.py b/tests/functional/diskusage/test_mila_diskusage.py index f96ef9aa..97ca6583 100644 --- a/tests/functional/diskusage/test_mila_diskusage.py +++ b/tests/functional/diskusage/test_mila_diskusage.py @@ -1,6 +1,8 @@ +import logging from pathlib import Path import pytest +from opentelemetry.trace import StatusCode import sarc.storage.mila from sarc.client.users.api import Credentials, User @@ -62,9 +64,10 @@ def mock_get_users(): "test_config", [{"clusters": {"mila": {"host": "mila"}}}], indirect=True ) def test_mila_fetch_diskusage_single( - test_config, monkeypatch, cli_main, file_regression + test_config, monkeypatch, cli_main, file_regression, caplog, captrace ): count = 0 + caplog.set_level(logging.INFO) def mock_get_report(*args): nonlocal count @@ -93,6 +96,17 @@ def mock_get_report(*args): # report = sarc.storage.mila.fetch_diskusage_report(cluster=test_config.clusters["mila"]) file_regression.check(data[0].json(exclude={"id": True}, indent=4)) + # check logging + assert caplog.text.__contains__("Acquiring mila storages report...") + + # check traces + traces = captrace.get_finished_spans() + assert len(traces) == 1 # only one cluster was acquired + assert traces[0].name == "cluster" + assert traces[0].attributes["cluster_name"] == "mila" + assert "Acquiring mila storages report..." in [e.name for e in traces[0].events] + assert traces[0].status.status_code == StatusCode.OK + @pytest.mark.freeze_time("2023-07-25") @pytest.mark.parametrize(