Make the tests work (hopefully)

Author: abergeron

sarc/cli/acquire/storages.py

@@ -22,29 +22,31 @@ def execute(self) -> int:
             print(f"Acquiring {cluster_name} storages report...")
 
             cluster = cfg.clusters[cluster_name]
-            diskusage = cluster.diskusage
-            if diskusage is None:
+            diskusage_configs = cluster.diskusage
+            if diskusage_configs is None:
                 continue
 
-            try:
-                scraper = get_diskusage_scraper(diskusage.name)
-            except KeyError as ke:
-                logger.exception(
-                    "Invalid or absent diskusage scraper name: %s",
-                    diskusage.name,
-                    exc_info=ke,
-                )
-                continue
-
-            disk_config = scraper.validate_config(diskusage.params)
-            data = scraper.get_diskusage_report(cluster.ssh, disk_config)
-            du = scraper.parse_diskusage_report(config, data)
-
-            if not self.dry:
-                collection = get_diskusage_collection()
-                collection.add(du)
-            else:
-                print("Document:")
-                print(du.model_dump_json(indent=2))
+            # Process each diskusage configuration
+            for diskusage_config in diskusage_configs:
+                try:
+                    scraper = get_diskusage_scraper(diskusage_config.name)
+                except KeyError as ke:
+                    logger.exception(
+                        "Invalid or absent diskusage scraper name: %s",
+                        diskusage_config.name,
+                        exc_info=ke,
+                    )
+                    continue
+
+                disk_config = scraper.validate_config(diskusage_config.params)
+                data = scraper.get_diskusage_report(cluster.ssh, disk_config)
+                du = scraper.parse_diskusage_report(disk_config, cluster_name, data)
+
+                if not self.dry:
+                    collection = get_diskusage_collection()
+                    collection.add(du)
+                else:
+                    print("Document:")
+                    print(du.model_dump_json(indent=2))
 
         return 0

sarc/config.py

@@ -42,7 +42,7 @@ class ConfigurationError(Exception):
 @dataclass
 class DiskUsageConfig:
     name: str
-    params: JSON = {}
+    params: JSON = field(default_factory=dict)
 
 
 @dataclass

sarc/storage/__init__.py

@@ -1 +1 @@
-from . import beegfs
+from . import beegfs, drac

sarc/storage/beegfs.py

@@ -4,7 +4,7 @@
 
 import json
 from dataclasses import dataclass
-from datetime import datetime
+from datetime import UTC, datetime
 from typing import cast
 
 from fabric import Connection
@@ -76,7 +76,7 @@ def parse_diskusage_report(
         return DiskUsage(
             cluster_name=cluster_name,
             groups=groups,
-            timestamp=datetime.now(),
+            timestamp=datetime.now(UTC),
         )
 
 

sarc/storage/diskusage.py

@@ -3,6 +3,7 @@
 from datetime import UTC, date, datetime
 from typing import Any
 
+from pydantic import field_validator
 from pydantic_mongo import AbstractRepository, PydanticObjectId
 
 from sarc.config import config
@@ -17,6 +18,19 @@ class DiskUsageDB(DiskUsage):
     # Database ID
     id: PydanticObjectId | None = None
 
+    @field_validator("timestamp", mode="before")
+    @classmethod
+    def _ensure_timestamp_utc(cls, value: str | datetime) -> datetime:
+        """Ensure timestamp has UTC timezone when reading from database."""
+        if isinstance(value, str):
+            return datetime.fromisoformat(value).replace(tzinfo=UTC)
+        elif isinstance(value, datetime):
+            if value.tzinfo is None:
+                # Convert naive datetime to UTC (from database storage)
+                return value.replace(tzinfo=UTC)
+            return value
+        return value
+
 
 class ClusterDiskUsageRepository(AbstractRepository[DiskUsageDB]):
     class Meta:
@@ -31,6 +45,7 @@ def add(self, disk_usage: DiskUsage):
             year=disk_usage.timestamp.year,
             month=disk_usage.timestamp.month,
             day=disk_usage.timestamp.day,
+            tzinfo=UTC,
         )
         disk_usage.timestamp = day_at_midnight
         document = self.to_document(DiskUsageDB(**disk_usage.model_dump()))

sarc/storage/drac.py

@@ -129,6 +129,11 @@ def _parse_body(L_lines: list[str]) -> list[DiskUsageGroup]:
             m = re.match(r"^\s*([\w\.]+)\s+(\d+)\s+([\d\.]+)\s(\w+)\s*", line)
             assert m, f"If this line doesn't match, we've got a problem.\n{line}"
             username = m.group(1)
+
+            # Skip the "Total" line as it's a summary, not a user
+            if username == "Total":
+                continue
+
             nbr_files = int(m.group(2))
             size = f"{m.group(3)} {m.group(4)}"
             LD_results.append(

tests/functional/diskusage/mila_reports/report_all.txt

@@ -0,0 +1 @@
+{"default": ["user1,1500000082,51046633472,107374182400,201276,1000000", "user2,1500000082,51046633472,107374182400,201276,1000000", "user3,1500000082,51046633472,107374182400,201276,1000000"]}

tests/functional/diskusage/mila_reports/report_single_user.txt renamed to tests/functional/diskusage/mila_reports/report_user1.txt

@@ -2,6 +2,4 @@
 Quota information for storage pool Default (ID: 1):
 
 name,id,size,hard,files,hard
-user1,1500000082,51046633472,107374182400,201276,1000000
 user2,1500000082,51046633472,107374182400,201276,1000000
-user3,1500000082,51046633472,107374182400,201276,1000000

tests/functional/diskusage/mila_reports/report_multiple_user.txt renamed to tests/functional/diskusage/mila_reports/report_user2.txt

@@ -0,0 +1,5 @@
+
+Quota information for storage pool Default (ID: 1):
+
+name,id,size,hard,files,hard
+user3,1500000082,51046633472,107374182400,201276,1000000