Do not pass around the queue object, but initialize one for the object

Author: alvarolopez

caso/messenger/ssm.py

@@ -68,19 +68,17 @@ def __init__(self):
         except Exception as err:
             LOG.error(f"Failed to create path {CONF.ssm.output_path} because {err}")
             raise err
+        self.queue = dirq.QueueSimple.QueueSimple(CONF.ssm.output_path)
 
-    def _push_message_cloud(
-        self, queue: dirq.QueueSimple.QueueSimple, entries: typing.List[str]
-    ):
+    def _push_message_cloud(self, entries: typing.List[str]):
         """Push a compute message, formatted following the CloudRecord."""
         message = f"APEL-cloud-message: v{self.version_cloud}\n"
         aux = "\n%%\n".join(entries)
         message += f"{aux}\n"
-        queue.add(message.encode("utf-8"))
+        self.queue.add(message.encode("utf-8"))
 
     def _push_message_json(
         self,
-        queue: dirq.QueueSimple.QueueSimple,
         entries: typing.List[str],
         msg_type: str,
         version: str,
@@ -91,25 +89,19 @@ def _push_message_json(
             "Version": version,
             "UsageRecords": [json.loads(r) for r in entries],
         }
-        queue.add(json.dumps(message))
+        self.queue.add(json.dumps(message))
 
-    def _push_message_ip(
-        self, queue: dirq.QueueSimple.QueueSimple, entries: typing.List[str]
-    ):
+    def _push_message_ip(self, entries: typing.List[str]):
         """Push an IP message."""
-        self._push_message_json(
-            queue, entries, "APEL Public IP message", self.version_ip
-        )
+        self._push_message_json(entries, "APEL Public IP message", self.version_ip)
 
-    def _push_message_accelerator(
-        self, queue: dirq.QueueSimple.QueueSimple, entries: typing.List[str]
-    ):
+    def _push_message_accelerator(self, entries: typing.List[str]):
         """Push an accelerator message."""
         self._push_message_json(
-            queue, entries, "APEL-accelerator-message", self.version_accelerator
+            entries, "APEL-accelerator-message", self.version_accelerator
         )
 
-    def _push_message_storage(self, queue, entries):
+    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)
@@ -136,31 +128,29 @@ def _push_message_storage(self, queue, entries):
             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
-        queue.add(ETree.tostring(root))
+        self.queue.add(ETree.tostring(root))
 
     def _push(self, entries_cloud, entries_ip, entries_accelerator, entries_storage):
         """Push all messages, dividing them into smaller chunks.
 
         This method gets lists of messages to be pushed in smaller chucks as per GGUS
         ticket 143436: https://ggus.eu/index.php?mode=ticket_info&ticket_id=143436
         """
-        queue = dirq.QueueSimple.QueueSimple(CONF.ssm.output_path)
-
         for i in range(0, len(entries_cloud), CONF.ssm.max_size):
             entries = entries_cloud[i : i + CONF.ssm.max_size]  # noqa(E203)
-            self._push_message_cloud(queue, entries)
+            self._push_message_cloud(entries)
 
         for i in range(0, len(entries_ip), CONF.ssm.max_size):
             entries = entries_ip[i : i + CONF.ssm.max_size]  # noqa(E203)
-            self._push_message_ip(queue, entries)
+            self._push_message_ip(entries)
 
         for i in range(0, len(entries_accelerator), CONF.ssm.max_size):
             entries = entries_accelerator[i : i + CONF.ssm.max_size]  # noqa(E203)
-            self._push_message_accelerator(queue, entries)
+            self._push_message_accelerator(entries)
 
         for i in range(0, len(entries_storage), CONF.ssm.max_size):
             entries = entries_storage[i : i + CONF.ssm.max_size]  # noqa(E203)
-            self._push_message_storage(queue, entries)
+            self._push_message_storage(entries)
 
     def push(self, records):
         """Push all records to SSM.

caso/tests/conftest.py

@@ -310,12 +310,40 @@ def expected_entries_cloud():
     return ssm_entries
 
 
+@pytest.fixture
+def expected_message_cloud():
+    """Get a fixture for a complete Cloud message."""
+    message = (
+        "APEL-cloud-message: v0.4\n"
+        "CloudComputeService: Fake Cloud Service\n"
+        f"CloudType: {cloud_type}\nCpuCount: 8\nCpuDuration: 3456000\n"
+        "Disk: 250\nEndTime: 1685044746\nFQAN: VO 1 FQAN\nGlobalUserName: User DN\n"
+        "ImageId: b39a8ed9-e15d-4b71-ada2-daf88efbac0a\n"
+        "LocalGroupId: 03b6a6c4-cf2b-48b9-82f1-69c52b9f30af\n"
+        "LocalUserId: a4519d7d-f60a-4908-9d63-7d9e17422188\nMachineName: VM Name 1\n"
+        "Memory: 16\nPublicIPCount: 7\nSiteName: TEST-Site\nStartTime: 1684612746\n"
+        "Status: started\nVMUUID: 721cf1db-0e0f-4c24-a5ea-cd75e0f303e8\n"
+        "WallDuration: 432000\n"
+        "%%"
+        "\nCloudComputeService: Fake Cloud Service\n"
+        f"CloudType: {cloud_type}\nCpuCount: 8\nCpuDuration: 3456000\n"
+        "Disk: 250\nEndTime: 1685044746\nFQAN: VO 2 FQAN\nGlobalUserName: User DN\n"
+        "ImageId: b39a8ed9-e15d-4b71-ada2-daf88efbac0a\n"
+        "LocalGroupId: 03b6a6c4-cf2b-48b9-82f1-69c52b9f30af\n"
+        "LocalUserId: a4519d7d-f60a-4908-9d63-7d9e17422188\nMachineName: VM Name 2\n"
+        "Memory: 16\nPublicIPCount: 7\nSiteName: TEST-Site\nStartTime: 1684526346\n"
+        "Status: completed\nVMUUID: a53738e1-13eb-4047-800c-067d14ce3d22\n"
+        "WallDuration: 432000\n"
+    )
+    return message.encode("utf-8")
+
+
 @pytest.fixture
 def expected_entries_ip():
     """Get a fixture for all IP entries."""
     ssm_entries = [
         '{"SiteName": "TEST-Site", '
-        '"CloudType": "caso/4.1.0.0rc1 (OpenStack)", '
+        f'"CloudType": "{cloud_type}", '
         '"CloudComputeService": "Fake Cloud Service", '
         '"uuid": "e3c5aeef-37b8-4332-ad9f-9d068f156dc2", '
         '"LocalUser": "a4519d7d-f60a-4908-9d63-7d9e17422188", '
@@ -326,7 +354,7 @@ def expected_entries_ip():
         '"IPVersion": 4, '
         '"IPCount": 10}',
         '{"SiteName": "TEST-Site", '
-        '"CloudType": "caso/4.1.0.0rc1 (OpenStack)", '
+        f'"CloudType": "{cloud_type}", '
         '"CloudComputeService": "Fake Cloud Service", '
         '"uuid": "5c50720e-a653-4d70-9b0e-d4388687fcbc", '
         '"LocalUser": "3391a44e-3728-478d-abde-b86c25356571", '
@@ -338,3 +366,35 @@ def expected_entries_ip():
         '"IPCount": 20}',
     ]
     return ssm_entries
+
+
+@pytest.fixture
+def expected_message_ip():
+    """Get a fixture for a complete IP message."""
+    message = (
+        '{"Type": "APEL Public IP message", "Version": "0.2", "UsageRecords": ['
+        '{"SiteName": "TEST-Site", '
+        f'"CloudType": "{cloud_type}", '
+        '"CloudComputeService": "Fake Cloud Service", '
+        '"uuid": "e3c5aeef-37b8-4332-ad9f-9d068f156dc2", '
+        '"LocalUser": "a4519d7d-f60a-4908-9d63-7d9e17422188", '
+        '"GlobalUserName": "User 1 DN", '
+        '"LocalGroup": "03b6a6c4-cf2b-48b9-82f1-69c52b9f30af", '
+        '"FQAN": "VO 1 FQAN", '
+        '"MeasurementTime": 1685044746, '
+        '"IPVersion": 4, '
+        '"IPCount": 10}, '
+        '{"SiteName": "TEST-Site", '
+        f'"CloudType": "{cloud_type}", '
+        '"CloudComputeService": "Fake Cloud Service", '
+        '"uuid": "5c50720e-a653-4d70-9b0e-d4388687fcbc", '
+        '"LocalUser": "3391a44e-3728-478d-abde-b86c25356571", '
+        '"GlobalUserName": "User 2 DN", '
+        '"LocalGroup": "2dae43c4-1889-4e63-b172-d4e99381e30a", '
+        '"FQAN": "VO 2 FQAN", '
+        '"MeasurementTime": 1685044746, '
+        '"IPVersion": 6, '
+        '"IPCount": 20}'
+        "]}"
+    )
+    return message

caso/tests/test_ssm.py

@@ -26,6 +26,7 @@ def test_empty_records_does_nothing(monkeypatch):
     """Test that empty records do nothing."""
     with monkeypatch.context() as m:
         m.setattr("caso.utils.makedirs", lambda x: None)
+        m.setattr("dirq.QueueSimple.QueueSimple", lambda x: None)
         messenger = ssm.SSMMessenger()
 
         assert messenger.push([]) is None
@@ -35,6 +36,7 @@ def test_weird_record_raises(monkeypatch):
     """Test that empty records do nothing."""
     with monkeypatch.context() as m:
         m.setattr("caso.utils.makedirs", lambda x: None)
+        m.setattr("dirq.QueueSimple.QueueSimple", lambda x: None)
         messenger = ssm.SSMMessenger()
 
         with pytest.raises(caso.exception.CasoError):
@@ -49,6 +51,7 @@ def mock_push(entries_cloud, entries_ip, entries_accelerator, entries_storage):
 
     with monkeypatch.context() as m:
         m.setattr("caso.utils.makedirs", lambda x: None)
+        m.setattr("dirq.QueueSimple.QueueSimple", lambda x: None)
         messenger = ssm.SSMMessenger()
 
         m.setattr(messenger, "_push", mock_push)
@@ -63,6 +66,7 @@ def mock_push(entries_cloud, entries_ip, entries_accelerator, entries_storage):
 
     with monkeypatch.context() as m:
         m.setattr("caso.utils.makedirs", lambda x: None)
+        m.setattr("dirq.QueueSimple.QueueSimple", lambda x: None)
         messenger = ssm.SSMMessenger()
 
         m.setattr(messenger, "_push", mock_push)
@@ -84,7 +88,49 @@ def mock_push(entries_cloud, entries_ip, entries_accelerator, entries_storage):
 
     with monkeypatch.context() as m:
         m.setattr("caso.utils.makedirs", lambda x: None)
+        m.setattr("dirq.QueueSimple.QueueSimple", lambda x: None)
         messenger = ssm.SSMMessenger()
 
         m.setattr(messenger, "_push", mock_push)
         messenger.push(cloud_record_list + ip_record_list)
+
+
+class _MockQueue:
+    @staticmethod
+    def add(message):
+        """Add a message to the fake queue."""
+        pass
+
+
+def test_complete_cloud_message(
+    monkeypatch, expected_entries_cloud, expected_message_cloud
+):
+    """Test a complete cloud message."""
+
+    def mock_add(message):
+        assert message == expected_message_cloud
+
+    with monkeypatch.context() as m:
+        m.setattr("caso.utils.makedirs", lambda x: None)
+        m.setattr("dirq.QueueSimple.QueueSimple", lambda x: _MockQueue())
+        messenger = ssm.SSMMessenger()
+
+        m.setattr(messenger.queue, "add", mock_add)
+        messenger._push_message_cloud(expected_entries_cloud)
+
+
+def test_complete_ip_message(monkeypatch, expected_entries_ip, expected_message_ip):
+    """Test a complete cloud message."""
+
+    def mock_add(message):
+        print(message)
+        print(expected_message_ip)
+        assert message == expected_message_ip
+
+    with monkeypatch.context() as m:
+        m.setattr("caso.utils.makedirs", lambda x: None)
+        m.setattr("dirq.QueueSimple.QueueSimple", lambda x: _MockQueue())
+        messenger = ssm.SSMMessenger()
+
+        m.setattr(messenger.queue, "add", mock_add)
+        messenger._push_message_ip(expected_entries_ip)