Sarc 319 log trace diskusage

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

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(