Skip to content

Commit 0be25b5

Browse files
committed
Implement accelerator records
1 parent c65dde3 commit 0be25b5

File tree

4 files changed

+224
-1
lines changed

4 files changed

+224
-1
lines changed

caso/record.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,7 @@ class AcceleratorRecord(_BaseRecord):
214214
215215
"""
216216

217-
version = "0.1"
217+
version: str = pydantic.Field("0.1", exclude=True)
218218

219219
uuid: m_uuid.UUID
220220

@@ -270,6 +270,8 @@ def map_fields(field: str) -> str:
270270
"benchmark": "Benchmark",
271271
"accelerator_type": "Type",
272272
"model": "Model",
273+
"cloud_type": "CloudType",
274+
"compute_service": "CloudComputeService",
273275
}
274276
return d.get(field, field)
275277

caso/tests/conftest.py

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,68 @@
173173
},
174174
]
175175

176+
valid_accelerator_records_fields = [
177+
dict(
178+
measurement_month=6,
179+
measurement_year=2022,
180+
uuid="99cf5d02-a573-46a1-b90d-0f7327126876",
181+
fqan="VO 1 FQAN",
182+
compute_service="Fake Cloud Service",
183+
site_name="TEST-Site",
184+
count=3,
185+
available_duration=5000,
186+
accelerator_type="GPU",
187+
user_dn="[email protected]",
188+
model="Foobar A200",
189+
),
190+
dict(
191+
measurement_month=2,
192+
measurement_year=2022,
193+
uuid="99cf5d02-a573-46a1-b90d-0f7327126876",
194+
fqan="VO 1 FQAN",
195+
compute_service="Fake Cloud Service",
196+
site_name="TEST-Site",
197+
count=30,
198+
available_duration=5000,
199+
accelerator_type="GPU",
200+
user_dn="[email protected]",
201+
model="Foobar A300",
202+
),
203+
]
204+
205+
valid_accelerator_records_dict = [
206+
{
207+
"AccUUID": "99cf5d02-a573-46a1-b90d-0f7327126876",
208+
"AssociatedRecordType": "cloud",
209+
"AvailableDuration": 5000,
210+
"Count": 3,
211+
"FQAN": "VO 1 FQAN",
212+
"GlobalUserName": "[email protected]",
213+
"MeasurementMonth": 6,
214+
"MeasurementYear": 2022,
215+
"Model": "Foobar A200",
216+
"SiteName": "TEST-Site",
217+
"Type": "GPU",
218+
"CloudType": cloud_type,
219+
"CloudComputeService": "Fake Cloud Service",
220+
},
221+
{
222+
"AccUUID": "99cf5d02-a573-46a1-b90d-0f7327126876",
223+
"AssociatedRecordType": "cloud",
224+
"AvailableDuration": 5000,
225+
"Count": 30,
226+
"FQAN": "VO 1 FQAN",
227+
"GlobalUserName": "[email protected]",
228+
"MeasurementMonth": 2,
229+
"MeasurementYear": 2022,
230+
"Model": "Foobar A300",
231+
"SiteName": "TEST-Site",
232+
"Type": "GPU",
233+
"CloudType": cloud_type,
234+
"CloudComputeService": "Fake Cloud Service",
235+
},
236+
]
237+
176238
# Cloud Record fixtures
177239

178240

@@ -260,6 +322,41 @@ def ip_record_list(ip_record, another_ip_record):
260322
return [ip_record, another_ip_record]
261323

262324

325+
# Accelerator records
326+
327+
328+
@pytest.fixture(scope="module")
329+
def accelerator_record():
330+
"""Get a fixture for the AcceleratorRecord."""
331+
record = caso.record.AcceleratorRecord(**valid_accelerator_records_fields[0])
332+
return record
333+
334+
335+
@pytest.fixture(scope="module")
336+
def another_accelerator_record():
337+
"""Get another fixture for the AcceleratorRecord."""
338+
record = caso.record.AcceleratorRecord(**valid_accelerator_records_fields[1])
339+
return record
340+
341+
342+
@pytest.fixture(scope="module")
343+
def valid_accelerator_record():
344+
"""Get a fixture for a valid record."""
345+
return valid_accelerator_records_dict[0]
346+
347+
348+
@pytest.fixture(scope="module")
349+
def valid_accelerator_records():
350+
"""Get a fixture for valid records as a dict."""
351+
return valid_accelerator_records_dict
352+
353+
354+
@pytest.fixture(scope="module")
355+
def accelerator_record_list(accelerator_record, another_accelerator_record):
356+
"""Get a fixture for a list of Accelerator records."""
357+
return [accelerator_record, another_accelerator_record]
358+
359+
263360
# SSM entries
264361

265362

@@ -398,3 +495,74 @@ def expected_message_ip():
398495
"]}"
399496
)
400497
return message
498+
499+
500+
@pytest.fixture
501+
def expected_entries_accelerator():
502+
"""Get a fixture for all accelerator entries."""
503+
ssm_entries = [
504+
'{"SiteName": "TEST-Site", '
505+
f'"CloudType": "{cloud_type}", '
506+
'"CloudComputeService": "Fake Cloud Service", '
507+
'"AccUUID": "99cf5d02-a573-46a1-b90d-0f7327126876", '
508+
'"GlobalUserName": "[email protected]", '
509+
'"FQAN": "VO 1 FQAN", '
510+
'"Count": 3, '
511+
'"AvailableDuration": 5000, '
512+
'"MeasurementMonth": 6, '
513+
'"MeasurementYear": 2022, '
514+
'"AssociatedRecordType": "cloud", '
515+
'"Type": "GPU", '
516+
'"Model": "Foobar A200"}',
517+
'{"SiteName": "TEST-Site", '
518+
f'"CloudType": "{cloud_type}", '
519+
'"CloudComputeService": "Fake Cloud Service", '
520+
'"AccUUID": "99cf5d02-a573-46a1-b90d-0f7327126876", '
521+
'"GlobalUserName": "[email protected]", '
522+
'"FQAN": "VO 1 FQAN", '
523+
'"Count": 30, '
524+
'"AvailableDuration": 5000, '
525+
'"MeasurementMonth": 2, '
526+
'"MeasurementYear": 2022, '
527+
'"AssociatedRecordType": "cloud", '
528+
'"Type": "GPU", '
529+
'"Model": "Foobar A300"}',
530+
]
531+
532+
return ssm_entries
533+
534+
535+
@pytest.fixture
536+
def expected_message_accelerator():
537+
"""Get a fixture for a complete Accelerator message."""
538+
message = (
539+
'{"Type": "APEL-accelerator-message", "Version": "0.1", "UsageRecords": ['
540+
'{"SiteName": "TEST-Site", '
541+
f'"CloudType": "{cloud_type}", '
542+
'"CloudComputeService": "Fake Cloud Service", '
543+
'"AccUUID": "99cf5d02-a573-46a1-b90d-0f7327126876", '
544+
'"GlobalUserName": "[email protected]", '
545+
'"FQAN": "VO 1 FQAN", '
546+
'"Count": 3, '
547+
'"AvailableDuration": 5000, '
548+
'"MeasurementMonth": 6, '
549+
'"MeasurementYear": 2022, '
550+
'"AssociatedRecordType": "cloud", '
551+
'"Type": "GPU", '
552+
'"Model": "Foobar A200"}, '
553+
'{"SiteName": "TEST-Site", '
554+
f'"CloudType": "{cloud_type}", '
555+
'"CloudComputeService": "Fake Cloud Service", '
556+
'"AccUUID": "99cf5d02-a573-46a1-b90d-0f7327126876", '
557+
'"GlobalUserName": "[email protected]", '
558+
'"FQAN": "VO 1 FQAN", '
559+
'"Count": 30, '
560+
'"AvailableDuration": 5000, '
561+
'"MeasurementMonth": 2, '
562+
'"MeasurementYear": 2022, '
563+
'"AssociatedRecordType": "cloud", '
564+
'"Type": "GPU", '
565+
'"Model": "Foobar A300"}'
566+
"]}"
567+
)
568+
return message

caso/tests/test_record.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,3 +105,17 @@ def test_ip_record_map_opts(ip_record, valid_ip_record):
105105
}
106106

107107
assert json.loads(ip_record.json(**opts)) == valid_ip_record
108+
109+
110+
def test_accelerator_record(accelerator_record):
111+
"""Test that an IP record is correctly generated."""
112+
assert isinstance(accelerator_record.available_duration, int)
113+
114+
115+
def test_accelerator_record_map_opts(accelerator_record, valid_accelerator_record):
116+
"""Test that an IP record is correctly generated."""
117+
opts = {
118+
"by_alias": True,
119+
"exclude_none": True,
120+
}
121+
assert json.loads(accelerator_record.json(**opts)) == valid_accelerator_record

caso/tests/test_ssm.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,26 @@ def mock_push(entries_cloud, entries_ip, entries_accelerator, entries_storage):
7373
messenger.push(ip_record_list)
7474

7575

76+
def test_accelerator_records_pushed(
77+
monkeypatch, accelerator_record_list, expected_entries_accelerator
78+
):
79+
"""Test that Accelerator records are correctly rendered."""
80+
81+
def mock_push(entries_cloud, entries_ip, entries_accelerator, entries_storage):
82+
import pprint
83+
84+
pprint.pprint(entries_accelerator)
85+
assert set(entries_accelerator) == set(expected_entries_accelerator)
86+
87+
with monkeypatch.context() as m:
88+
m.setattr("caso.utils.makedirs", lambda x: None)
89+
m.setattr("dirq.QueueSimple.QueueSimple", lambda x: None)
90+
messenger = ssm.SSMMessenger()
91+
92+
m.setattr(messenger, "_push", mock_push)
93+
messenger.push(accelerator_record_list)
94+
95+
7696
def test_cloud_ip_records_pushed(
7797
monkeypatch,
7898
cloud_record_list,
@@ -134,3 +154,22 @@ def mock_add(message):
134154

135155
m.setattr(messenger.queue, "add", mock_add)
136156
messenger._push_message_ip(expected_entries_ip)
157+
158+
159+
def test_complete_accelerator_message(
160+
monkeypatch, expected_entries_accelerator, expected_message_accelerator
161+
):
162+
"""Test a complete cloud message."""
163+
164+
def mock_add(message):
165+
print(message)
166+
print(expected_message_accelerator)
167+
assert message == expected_message_accelerator
168+
169+
with monkeypatch.context() as m:
170+
m.setattr("caso.utils.makedirs", lambda x: None)
171+
m.setattr("dirq.QueueSimple.QueueSimple", lambda x: _MockQueue())
172+
messenger = ssm.SSMMessenger()
173+
174+
m.setattr(messenger.queue, "add", mock_add)
175+
messenger._push_message_accelerator(expected_entries_accelerator)

0 commit comments

Comments
 (0)