Records must implement an SSM rendering by themselves

alvarolopez · alvarolopez · commit fb37c0ad695a · 2023-06-07T10:01:49.000+02:00
Instead of rendering the individual records in the SSM messenger, the
records classes must now implement a `ssm_message` method that provides
an export of the mode, with correct serialized values for each field,
that the SSM messenger will append to a SSM message.
diff --git a/caso/messenger/ssm.py b/caso/messenger/ssm.py
@@ -103,31 +103,14 @@ def _push_message_accelerator(self, entries: typing.List[str]):
 
     def _push_message_storage(self, entries):
         """Push a storage message."""
-        ns = {"xmlns:sr": "http://eu-emi.eu/namespaces/2011/02/storagerecord"}
-        root = ETree.Element("sr:StorageUsageRecords", attrib=ns)
+        ETree.register_namespace(
+            "sr", "http://eu-emi.eu/namespaces/2011/02/storagerecord"
+        )
+        root = ETree.Element("sr:StorageUsageRecords")
         for record in entries:
-            sr = ETree.SubElement(root, "sr:StorageUsageRecord")
-            ETree.SubElement(
-                sr,
-                "sr:RecordIdentity",
-                attrib={
-                    "sr:createTime": record.measure_time.isoformat(),
-                    "sr:recordId": str(record.uuid),
-                },
-            )
-            ETree.SubElement(sr, "sr:StorageSystem").text = record.compute_service
-            ETree.SubElement(sr, "sr:Site").text = record.site_name
-            subject = ETree.SubElement(sr, "sr:SubjectIdentity")
-            ETree.SubElement(subject, "sr:LocalUser").text = record.user_id
-            ETree.SubElement(subject, "sr:LocalGroup").text = record.group_id
-            if record.user_dn:
-                ETree.SubElement(subject, "sr:UserIdentity").text = record.user_dn
-            if record.fqan:
-                ETree.SubElement(subject, "sr:Group").text = record.fqan
-            ETree.SubElement(sr, "sr:StartTime").text = record.start_time.isoformat()
-            ETree.SubElement(sr, "sr:EndTime").text = record.measure_time.isoformat()
-            capacity = str(int(record.capacity * 1073741824))  # 1 GiB = 2^30
-            ETree.SubElement(sr, "sr:ResourceCapacityUsed").text = capacity
+            # We are not parsing XML so this is safe
+            sr = ETree.fromstring(record)  # nosec
+            root.append(sr)
         self.queue.add(ETree.tostring(root))
 
     def _push(self, entries_cloud, entries_ip, entries_accelerator, entries_storage):
@@ -171,28 +154,15 @@ def push(self, records):
         entries_ip = []
         entries_accelerator = []
         entries_storage = []
-        opts = {
-            "by_alias": True,
-            "exclude_none": True,
-        }
         for record in records:
             if isinstance(record, caso.record.CloudRecord):
-                aux = []
-                # NOTE(aloga): do not iter over the dictionary returned by record.dict()
-                # as this is just a dictionary representation of the object, where no
-                # serialization is done. In order to get objects correctly serialized
-                # we need to convert to JSON, then reload the model
-                serialized_record = json.loads(record.json(**opts))
-                for k, v in serialized_record.items():
-                    aux.append(f"{k}: {v}")
-                aux.sort()
-                entries_cloud.append("\n".join(aux))
+                entries_cloud.append(record.ssm_message())
             elif isinstance(record, caso.record.IPRecord):
-                entries_ip.append(record.json(**opts))
+                entries_ip.append(record.ssm_message())
             elif isinstance(record, caso.record.AcceleratorRecord):
-                entries_accelerator.append(record.json(**opts))
+                entries_accelerator.append(record.ssm_message())
             elif isinstance(record, caso.record.StorageRecord):
-                entries_storage.append(record)
+                entries_storage.append(record.ssm_message())
             else:
                 raise caso.exception.CasoError("Unexpected record format!")
 
diff --git a/caso/record.py b/caso/record.py
@@ -16,11 +16,16 @@
 
 """Module containing all the cloud accounting records."""
 
+import abc
 import datetime
 import enum
+import json
 import typing
 import uuid as m_uuid
 
+# We are not parsing XML so this is safe
+import xml.etree.ElementTree as ETree  # nosec
+
 import pydantic
 
 import caso
@@ -29,7 +34,7 @@
 LOG = log.getLogger(__name__)
 
 
-class _BaseRecord(pydantic.BaseModel):
+class _BaseRecord(pydantic.BaseModel, abc.ABC):
     """This is the base cASO record object."""
 
     version: str = pydantic.Field(..., exclude=True)
@@ -38,6 +43,11 @@ class _BaseRecord(pydantic.BaseModel):
     cloud_type: str = caso.user_agent
     compute_service: str
 
+    @abc.abstractmethod
+    def ssm_message(self):
+        """Render record as the expected SSM message."""
+        raise NotImplementedError("Method not implemented")
+
 
 class _ValidCloudStatus(str, enum.Enum):
     started = "started"
@@ -117,6 +127,21 @@ def set_cpu_duration(self, value: int):
         """Set the CPU duration for the record."""
         self._cpu_duration = value
 
+    def ssm_message(self):
+        """Render record as the expected SSM message."""
+        opts = {
+            "by_alias": True,
+            "exclude_none": True,
+        }
+        # NOTE(aloga): do not iter over the dictionary returned by record.dict() as this
+        # is just a dictionary representation of the object, where no serialization is
+        # done. In order to get objects correctly serialized we need to convert to JSON,
+        # then reload the model
+        serialized_record = json.loads(self.json(**opts))
+        aux = [f"{k}: {v}" for k, v in serialized_record.items()]
+        aux.sort()
+        return "\n".join(aux)
+
     class Config:
         """Config class for Pydantic."""
 
@@ -180,6 +205,14 @@ class IPRecord(_BaseRecord):
     ip_version: int
     public_ip_count: int
 
+    def ssm_message(self):
+        """Render record as the expected SSM message."""
+        opts = {
+            "by_alias": True,
+            "exclude_none": True,
+        }
+        return self.json(**opts)
+
     class Config:
         """Config class for Pydantic."""
 
@@ -248,6 +281,14 @@ def set_active_duration(self, value: int):
         """Set the active duration for the record."""
         self._active_duration = value
 
+    def ssm_message(self):
+        """Render record as the expected SSM message."""
+        opts = {
+            "by_alias": True,
+            "exclude_none": True,
+        }
+        return self.json(**opts)
+
     class Config:
         """Config class for Pydantic."""
 
@@ -318,6 +359,33 @@ def _validate_attached_duration(cls, value):
             return value
         return 0
 
+    def ssm_message(self):
+        """Render record as the expected SSM message."""
+        ns = {"xmlns:sr": "http://eu-emi.eu/namespaces/2011/02/storagerecord"}
+        sr = ETree.Element("sr:StorageUsageRecord", attrib=ns)
+        ETree.SubElement(
+            sr,
+            "sr:RecordIdentity",
+            attrib={
+                "sr:createTime": self.measure_time.isoformat(),
+                "sr:recordId": str(self.uuid),
+            },
+        )
+        ETree.SubElement(sr, "sr:StorageSystem").text = self.compute_service
+        ETree.SubElement(sr, "sr:Site").text = self.site_name
+        subject = ETree.SubElement(sr, "sr:SubjectIdentity")
+        ETree.SubElement(subject, "sr:LocalUser").text = self.user_id
+        ETree.SubElement(subject, "sr:LocalGroup").text = self.group_id
+        if self.user_dn:
+            ETree.SubElement(subject, "sr:UserIdentity").text = self.user_dn
+        if self.fqan:
+            ETree.SubElement(subject, "sr:Group").text = self.fqan
+        ETree.SubElement(sr, "sr:StartTime").text = self.start_time.isoformat()
+        ETree.SubElement(sr, "sr:EndTime").text = self.measure_time.isoformat()
+        capacity = str(int(self.capacity * 1073741824))  # 1 GiB = 2^30
+        ETree.SubElement(sr, "sr:ResourceCapacityUsed").text = capacity
+        return ETree.tostring(sr)
+
     class Config:
         """Config class for Pydantic."""
 
diff --git a/caso/tests/conftest.py b/caso/tests/conftest.py
@@ -681,9 +681,38 @@ def expected_message_accelerator() -> str:
 
 
 @pytest.fixture
-def expected_entries_storage(storage_record_list) -> list[caso.record.StorageRecord]:
+def expected_entries_storage() -> typing.List[str]:
     """Get a fixture for all Storage entries."""
-    ssm_entries = storage_record_list.copy()
+    ssm_entries = [
+        '<sr:StorageUsageRecord xmlns:sr="http://eu-emi.eu/namespaces/2011/02/storagerecord">'  # noqa
+        '<sr:RecordIdentity sr:createTime="2023-05-25T21:59:06+00:00" sr:recordId="99cf5d02-a573-46a1-b90d-0f7327126876" />'  # noqa
+        "<sr:StorageSystem>Fake Cloud Service</sr:StorageSystem>"
+        "<sr:Site>TEST-Site</sr:Site>"
+        "<sr:SubjectIdentity>"
+        "<sr:LocalUser>63296dcd-b652-4039-b274-aaa70f9d57e5</sr:LocalUser>"
+        "<sr:LocalGroup>313c6f62-e05f-4ec7-b0f2-256612db18f5</sr:LocalGroup>"
+        "<sr:UserIdentity>d4e547e6f298fe34389@foobar.eu</sr:UserIdentity>"
+        "<sr:Group>VO 1 FQAN</sr:Group>"
+        "</sr:SubjectIdentity>"
+        "<sr:StartTime>2023-05-20T21:59:06+00:00</sr:StartTime>"
+        "<sr:EndTime>2023-05-25T21:59:06+00:00</sr:EndTime>"
+        "<sr:ResourceCapacityUsed>345876451382054092800</sr:ResourceCapacityUsed>"
+        "</sr:StorageUsageRecord>",
+        '<sr:StorageUsageRecord xmlns:sr="http://eu-emi.eu/namespaces/2011/02/storagerecord">'  # noqa
+        '<sr:RecordIdentity sr:createTime="2023-05-25T21:59:06+00:00" sr:recordId="99cf5d02-a573-46a1-b90d-0f7327126876" />'  # noqa
+        "<sr:StorageSystem>Fake Cloud Service</sr:StorageSystem>"
+        "<sr:Site>TEST-Site</sr:Site>"
+        "<sr:SubjectIdentity>"
+        "<sr:LocalUser>63296dcd-b652-4039-b274-aaa70f9d57e5</sr:LocalUser>"
+        "<sr:LocalGroup>313c6f62-e05f-4ec7-b0f2-256612db18f5</sr:LocalGroup>"
+        "<sr:UserIdentity>d4e547e6f298fe34389@foobar.eu</sr:UserIdentity>"
+        "<sr:Group>VO 2 FQAN</sr:Group>"
+        "</sr:SubjectIdentity>"
+        "<sr:StartTime>2023-05-19T21:59:06+00:00</sr:StartTime>"
+        "<sr:EndTime>2023-05-25T21:59:06+00:00</sr:EndTime>"
+        "<sr:ResourceCapacityUsed>131128086582054092800</sr:ResourceCapacityUsed>"
+        "</sr:StorageUsageRecord>",
+    ]
     return ssm_entries
 
 
@@ -722,4 +751,4 @@ def expected_message_storage() -> str:
         "</sr:StorageUsageRecord>"
         "</sr:StorageUsageRecords>"
     )
-    return message.encode("utf-8")
+    return message
diff --git a/caso/tests/test_ssm.py b/caso/tests/test_ssm.py
@@ -96,7 +96,9 @@ def test_storage_records_pushed(
     """Test that Storage records are correctly rendered."""
 
     def mock_push(entries_cloud, entries_ip, entries_accelerator, entries_storage):
-        assert entries_storage == expected_entries_storage
+        assert set([s.decode() for s in entries_storage]) == set(
+            expected_entries_storage
+        )
 
     with monkeypatch.context() as m:
         m.setattr("caso.utils.makedirs", lambda x: None)
@@ -191,7 +193,7 @@ def test_complete_storage_message(
     """Test a complete cloud message."""
 
     def mock_add(message):
-        assert message == expected_message_storage
+        assert message.decode() == expected_message_storage
 
     with monkeypatch.context() as m:
         m.setattr("caso.utils.makedirs", lambda x: None)
