Skip to content

Commit fb37c0a

Browse files
alvarolopezalvarolopez
committed
Records must implement an SSM rendering by themselves
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.
1 parent b9457a6 commit fb37c0a

File tree

4 files changed

+116
-47
lines changed

4 files changed

+116
-47
lines changed

‎caso/messenger/ssm.py‎

Lines changed: 11 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -103,31 +103,14 @@ def _push_message_accelerator(self, entries: typing.List[str]):
103103

104104
def _push_message_storage(self, entries):
105105
"""Push a storage message."""
106-
ns = {"xmlns:sr": "http://eu-emi.eu/namespaces/2011/02/storagerecord"}
107-
root = ETree.Element("sr:StorageUsageRecords", attrib=ns)
106+
ETree.register_namespace(
107+
"sr", "http://eu-emi.eu/namespaces/2011/02/storagerecord"
108+
)
109+
root = ETree.Element("sr:StorageUsageRecords")
108110
for record in entries:
109-
sr = ETree.SubElement(root, "sr:StorageUsageRecord")
110-
ETree.SubElement(
111-
sr,
112-
"sr:RecordIdentity",
113-
attrib={
114-
"sr:createTime": record.measure_time.isoformat(),
115-
"sr:recordId": str(record.uuid),
116-
},
117-
)
118-
ETree.SubElement(sr, "sr:StorageSystem").text = record.compute_service
119-
ETree.SubElement(sr, "sr:Site").text = record.site_name
120-
subject = ETree.SubElement(sr, "sr:SubjectIdentity")
121-
ETree.SubElement(subject, "sr:LocalUser").text = record.user_id
122-
ETree.SubElement(subject, "sr:LocalGroup").text = record.group_id
123-
if record.user_dn:
124-
ETree.SubElement(subject, "sr:UserIdentity").text = record.user_dn
125-
if record.fqan:
126-
ETree.SubElement(subject, "sr:Group").text = record.fqan
127-
ETree.SubElement(sr, "sr:StartTime").text = record.start_time.isoformat()
128-
ETree.SubElement(sr, "sr:EndTime").text = record.measure_time.isoformat()
129-
capacity = str(int(record.capacity * 1073741824)) # 1 GiB = 2^30
130-
ETree.SubElement(sr, "sr:ResourceCapacityUsed").text = capacity
111+
# We are not parsing XML so this is safe
112+
sr = ETree.fromstring(record) # nosec
113+
root.append(sr)
131114
self.queue.add(ETree.tostring(root))
132115

133116
def _push(self, entries_cloud, entries_ip, entries_accelerator, entries_storage):
@@ -171,28 +154,15 @@ def push(self, records):
171154
entries_ip = []
172155
entries_accelerator = []
173156
entries_storage = []
174-
opts = {
175-
"by_alias": True,
176-
"exclude_none": True,
177-
}
178157
for record in records:
179158
if isinstance(record, caso.record.CloudRecord):
180-
aux = []
181-
# NOTE(aloga): do not iter over the dictionary returned by record.dict()
182-
# as this is just a dictionary representation of the object, where no
183-
# serialization is done. In order to get objects correctly serialized
184-
# we need to convert to JSON, then reload the model
185-
serialized_record = json.loads(record.json(**opts))
186-
for k, v in serialized_record.items():
187-
aux.append(f"{k}: {v}")
188-
aux.sort()
189-
entries_cloud.append("\n".join(aux))
159+
entries_cloud.append(record.ssm_message())
190160
elif isinstance(record, caso.record.IPRecord):
191-
entries_ip.append(record.json(**opts))
161+
entries_ip.append(record.ssm_message())
192162
elif isinstance(record, caso.record.AcceleratorRecord):
193-
entries_accelerator.append(record.json(**opts))
163+
entries_accelerator.append(record.ssm_message())
194164
elif isinstance(record, caso.record.StorageRecord):
195-
entries_storage.append(record)
165+
entries_storage.append(record.ssm_message())
196166
else:
197167
raise caso.exception.CasoError("Unexpected record format!")
198168

‎caso/record.py‎

Lines changed: 69 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,11 +16,16 @@
1616

1717
"""Module containing all the cloud accounting records."""
1818

19+
import abc
1920
import datetime
2021
import enum
22+
import json
2123
import typing
2224
import uuid as m_uuid
2325

26+
# We are not parsing XML so this is safe
27+
import xml.etree.ElementTree as ETree # nosec
28+
2429
import pydantic
2530

2631
import caso
@@ -29,7 +34,7 @@
2934
LOG = log.getLogger(__name__)
3035

3136

32-
class _BaseRecord(pydantic.BaseModel):
37+
class _BaseRecord(pydantic.BaseModel, abc.ABC):
3338
"""This is the base cASO record object."""
3439

3540
version: str = pydantic.Field(..., exclude=True)
@@ -38,6 +43,11 @@ class _BaseRecord(pydantic.BaseModel):
3843
cloud_type: str = caso.user_agent
3944
compute_service: str
4045

46+
@abc.abstractmethod
47+
def ssm_message(self):
48+
"""Render record as the expected SSM message."""
49+
raise NotImplementedError("Method not implemented")
50+
4151

4252
class _ValidCloudStatus(str, enum.Enum):
4353
started = "started"
@@ -117,6 +127,21 @@ def set_cpu_duration(self, value: int):
117127
"""Set the CPU duration for the record."""
118128
self._cpu_duration = value
119129

130+
def ssm_message(self):
131+
"""Render record as the expected SSM message."""
132+
opts = {
133+
"by_alias": True,
134+
"exclude_none": True,
135+
}
136+
# NOTE(aloga): do not iter over the dictionary returned by record.dict() as this
137+
# is just a dictionary representation of the object, where no serialization is
138+
# done. In order to get objects correctly serialized we need to convert to JSON,
139+
# then reload the model
140+
serialized_record = json.loads(self.json(**opts))
141+
aux = [f"{k}: {v}" for k, v in serialized_record.items()]
142+
aux.sort()
143+
return "\n".join(aux)
144+
120145
class Config:
121146
"""Config class for Pydantic."""
122147

@@ -180,6 +205,14 @@ class IPRecord(_BaseRecord):
180205
ip_version: int
181206
public_ip_count: int
182207

208+
def ssm_message(self):
209+
"""Render record as the expected SSM message."""
210+
opts = {
211+
"by_alias": True,
212+
"exclude_none": True,
213+
}
214+
return self.json(**opts)
215+
183216
class Config:
184217
"""Config class for Pydantic."""
185218

@@ -248,6 +281,14 @@ def set_active_duration(self, value: int):
248281
"""Set the active duration for the record."""
249282
self._active_duration = value
250283

284+
def ssm_message(self):
285+
"""Render record as the expected SSM message."""
286+
opts = {
287+
"by_alias": True,
288+
"exclude_none": True,
289+
}
290+
return self.json(**opts)
291+
251292
class Config:
252293
"""Config class for Pydantic."""
253294

@@ -318,6 +359,33 @@ def _validate_attached_duration(cls, value):
318359
return value
319360
return 0
320361

362+
def ssm_message(self):
363+
"""Render record as the expected SSM message."""
364+
ns = {"xmlns:sr": "http://eu-emi.eu/namespaces/2011/02/storagerecord"}
365+
sr = ETree.Element("sr:StorageUsageRecord", attrib=ns)
366+
ETree.SubElement(
367+
sr,
368+
"sr:RecordIdentity",
369+
attrib={
370+
"sr:createTime": self.measure_time.isoformat(),
371+
"sr:recordId": str(self.uuid),
372+
},
373+
)
374+
ETree.SubElement(sr, "sr:StorageSystem").text = self.compute_service
375+
ETree.SubElement(sr, "sr:Site").text = self.site_name
376+
subject = ETree.SubElement(sr, "sr:SubjectIdentity")
377+
ETree.SubElement(subject, "sr:LocalUser").text = self.user_id
378+
ETree.SubElement(subject, "sr:LocalGroup").text = self.group_id
379+
if self.user_dn:
380+
ETree.SubElement(subject, "sr:UserIdentity").text = self.user_dn
381+
if self.fqan:
382+
ETree.SubElement(subject, "sr:Group").text = self.fqan
383+
ETree.SubElement(sr, "sr:StartTime").text = self.start_time.isoformat()
384+
ETree.SubElement(sr, "sr:EndTime").text = self.measure_time.isoformat()
385+
capacity = str(int(self.capacity * 1073741824)) # 1 GiB = 2^30
386+
ETree.SubElement(sr, "sr:ResourceCapacityUsed").text = capacity
387+
return ETree.tostring(sr)
388+
321389
class Config:
322390
"""Config class for Pydantic."""
323391

‎caso/tests/conftest.py‎

Lines changed: 32 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -681,9 +681,38 @@ def expected_message_accelerator() -> str:
681681

682682

683683
@pytest.fixture
684-
def expected_entries_storage(storage_record_list) -> list[caso.record.StorageRecord]:
684+
def expected_entries_storage() -> typing.List[str]:
685685
"""Get a fixture for all Storage entries."""
686-
ssm_entries = storage_record_list.copy()
686+
ssm_entries = [
687+
'<sr:StorageUsageRecord xmlns:sr="http://eu-emi.eu/namespaces/2011/02/storagerecord">' # noqa
688+
'<sr:RecordIdentity sr:createTime="2023-05-25T21:59:06+00:00" sr:recordId="99cf5d02-a573-46a1-b90d-0f7327126876" />' # noqa
689+
"<sr:StorageSystem>Fake Cloud Service</sr:StorageSystem>"
690+
"<sr:Site>TEST-Site</sr:Site>"
691+
"<sr:SubjectIdentity>"
692+
"<sr:LocalUser>63296dcd-b652-4039-b274-aaa70f9d57e5</sr:LocalUser>"
693+
"<sr:LocalGroup>313c6f62-e05f-4ec7-b0f2-256612db18f5</sr:LocalGroup>"
694+
"<sr:UserIdentity>[email protected]</sr:UserIdentity>"
695+
"<sr:Group>VO 1 FQAN</sr:Group>"
696+
"</sr:SubjectIdentity>"
697+
"<sr:StartTime>2023-05-20T21:59:06+00:00</sr:StartTime>"
698+
"<sr:EndTime>2023-05-25T21:59:06+00:00</sr:EndTime>"
699+
"<sr:ResourceCapacityUsed>345876451382054092800</sr:ResourceCapacityUsed>"
700+
"</sr:StorageUsageRecord>",
701+
'<sr:StorageUsageRecord xmlns:sr="http://eu-emi.eu/namespaces/2011/02/storagerecord">' # noqa
702+
'<sr:RecordIdentity sr:createTime="2023-05-25T21:59:06+00:00" sr:recordId="99cf5d02-a573-46a1-b90d-0f7327126876" />' # noqa
703+
"<sr:StorageSystem>Fake Cloud Service</sr:StorageSystem>"
704+
"<sr:Site>TEST-Site</sr:Site>"
705+
"<sr:SubjectIdentity>"
706+
"<sr:LocalUser>63296dcd-b652-4039-b274-aaa70f9d57e5</sr:LocalUser>"
707+
"<sr:LocalGroup>313c6f62-e05f-4ec7-b0f2-256612db18f5</sr:LocalGroup>"
708+
"<sr:UserIdentity>[email protected]</sr:UserIdentity>"
709+
"<sr:Group>VO 2 FQAN</sr:Group>"
710+
"</sr:SubjectIdentity>"
711+
"<sr:StartTime>2023-05-19T21:59:06+00:00</sr:StartTime>"
712+
"<sr:EndTime>2023-05-25T21:59:06+00:00</sr:EndTime>"
713+
"<sr:ResourceCapacityUsed>131128086582054092800</sr:ResourceCapacityUsed>"
714+
"</sr:StorageUsageRecord>",
715+
]
687716
return ssm_entries
688717

689718

@@ -722,4 +751,4 @@ def expected_message_storage() -> str:
722751
"</sr:StorageUsageRecord>"
723752
"</sr:StorageUsageRecords>"
724753
)
725-
return message.encode("utf-8")
754+
return message

‎caso/tests/test_ssm.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,9 @@ def test_storage_records_pushed(
9696
"""Test that Storage records are correctly rendered."""
9797

9898
def mock_push(entries_cloud, entries_ip, entries_accelerator, entries_storage):
99-
assert entries_storage == expected_entries_storage
99+
assert set([s.decode() for s in entries_storage]) == set(
100+
expected_entries_storage
101+
)
100102

101103
with monkeypatch.context() as m:
102104
m.setattr("caso.utils.makedirs", lambda x: None)
@@ -191,7 +193,7 @@ def test_complete_storage_message(
191193
"""Test a complete cloud message."""
192194

193195
def mock_add(message):
194-
assert message == expected_message_storage
196+
assert message.decode() == expected_message_storage
195197

196198
with monkeypatch.context() as m:
197199
m.setattr("caso.utils.makedirs", lambda x: None)

0 commit comments

Comments
 (0)