Auditor accounting plugin tests

tardis/plugins/auditor.py

@@ -4,7 +4,7 @@
 from ..utilities.attributedict import AttributeDict
 from ..resources.dronestates import AvailableState, DownState
 
-from pyauditor import AuditorClientBuilder, Record, Component, Score
+import pyauditor
 
 import logging
 import pytz
@@ -33,16 +33,16 @@ def __init__(self):
                     self._resources[site.name][machine_type][r] = getattr(
                         config, site.name
                     ).MachineMetaData[machine_type][r]
-                    self._components[site.name][machine_type][r] = config_auditor[
-                        "components"
-                    ][machine_type].get(r, {})
+                    self._components[site.name][machine_type][r] = getattr(
+                        config_auditor.components, machine_type
+                    ).get(r, {})
 
         self._user = config_auditor.user if config_auditor.user else "tardis"
         self._group = config_auditor.group if config_auditor.group else "tardis"
-        auditor_timeout = config_auditor.get("timeout", 30)
+        auditor_timeout = getattr(config_auditor, "timeout", 30)
         self._local_timezone = get_localzone()
         self._client = (
-            AuditorClientBuilder()
+            pyauditor.AuditorClientBuilder()
             .address(config_auditor.host, config_auditor.port)
             .timeout(auditor_timeout)
             .build()
@@ -76,7 +76,7 @@ async def notify(self, state: State, resource_attributes: AttributeDict) -> None
             await self._client.update(record)
 
     def construct_record(self, resource_attributes: AttributeDict):
-        record = Record(
+        record = pyauditor.Record(
             resource_attributes["drone_uuid"],
             resource_attributes["site_name"],
             self._user,
@@ -89,11 +89,11 @@ def construct_record(self, resource_attributes: AttributeDict):
         for (resource, amount) in self._resources[resource_attributes["site_name"]][
             resource_attributes["machine_type"]
         ].items():
-            component = Component(resource, amount)
+            component = pyauditor.Component(resource, amount)
             for score_name, score_val in self._components[
                 resource_attributes["site_name"]
             ][resource_attributes["machine_type"]][resource].items():
-                component = component.with_score(Score(score_name, score_val))
+                component = component.with_score(pyauditor.Score(score_name, score_val))
 
             record = record.with_component(component)
 

tests/plugins_t/test_auditor.py

@@ -0,0 +1,118 @@
+from tardis.plugins.auditor import Auditor
+from tardis.utilities.attributedict import AttributeDict
+from tardis.interfaces.siteadapter import ResourceStatus
+from tardis.resources.dronestates import AvailableState, DownState
+
+from datetime import datetime
+from unittest import TestCase
+from unittest.mock import patch
+
+from ..utilities.utilities import async_return
+from ..utilities.utilities import run_async
+
+
+class TestAuditor(TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.mock_config_patcher = patch("tardis.plugins.auditor.Configuration")
+        cls.mock_auditorclientbuilder_patcher = patch(
+            "tardis.plugins.auditor.pyauditor.AuditorClientBuilder"
+        )
+
+        cls.mock_config = cls.mock_config_patcher.start()
+        cls.mock_auditorclientbuilder = cls.mock_auditorclientbuilder_patcher.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.mock_config_patcher.stop()
+        cls.mock_auditorclientbuilder_patcher.stop()
+
+    def setUp(self):
+        self.address = "127.0.0.1"
+        self.port = 8000
+        self.timeout = 20
+        self.user = "user-1"
+        self.group = "group-1"
+        self.site = "testsite"
+        self.cores = 12
+        self.memory = 100
+        self.drone_uuid = "test-drone"
+        self.machine_type = "test_machine_type"
+        config = self.mock_config.return_value
+        config.Plugins.Auditor.host = self.address
+        config.Plugins.Auditor.port = self.port
+        config.Plugins.Auditor.timeout = self.timeout
+        config.Plugins.Auditor.user = self.user
+        config.Plugins.Auditor.group = self.group
+        config.Plugins.Auditor.components.test_machine_type = AttributeDict(
+            Cores=AttributeDict(HEPSPEC=1.2, BENCHMARK=3.0),
+            Memory=AttributeDict(BLUBB=1.4),
+        )
+        config.Sites = [AttributeDict(name=self.site)]
+        config.testsite.MachineTypes = [self.machine_type]
+        config.testsite.MachineMetaData = AttributeDict(
+            test_machine_type=AttributeDict(Cores=self.cores, Memory=self.memory)
+        )
+
+        self.test_param = AttributeDict(
+            site_name=self.site,
+            machine_type=self.machine_type,
+            created=datetime.now(),
+            updated=datetime.now(),
+            resource_status=ResourceStatus.Booting,
+            drone_uuid=self.drone_uuid,
+        )
+
+        builder = self.mock_auditorclientbuilder.return_value
+        builder = builder.address.return_value
+        builder = builder.timeout.return_value
+        self.client = builder.build.return_value
+        self.client.add.return_value = async_return()
+        self.client.update.return_value = async_return()
+
+        self.plugin = Auditor()
+
+    def test_notify(self):
+        self.mock_auditorclientbuilder.return_value.address.assert_called_with(
+            self.address,
+            self.port,
+        )
+        self.mock_auditorclientbuilder.return_value.address.return_value.timeout.assert_called_with(
+            self.timeout,
+        )
+        run_async(
+            self.plugin.notify,
+            state=AvailableState(),
+            resource_attributes=self.test_param,
+        )
+        self.assertEqual(self.client.add.call_count, 1)
+        self.assertEqual(self.client.update.call_count, 0)
+
+        run_async(
+            self.plugin.notify,
+            state=DownState(),
+            resource_attributes=self.test_param,
+        )
+        self.assertEqual(self.client.add.call_count, 1)
+        self.assertEqual(self.client.update.call_count, 1)
+
+    def test_construct_record(self):
+        record = self.plugin.construct_record(resource_attributes=self.test_param)
+
+        self.assertEqual(record.record_id, self.drone_uuid)
+        self.assertEqual(record.site_id, self.site)
+        self.assertEqual(record.user_id, self.user)
+        self.assertEqual(record.group_id, self.group)
+        self.assertEqual(len(record.components), 2)
+        self.assertEqual(record.components[0].name, "Cores")
+        self.assertEqual(record.components[0].amount, 12)
+        self.assertEqual(len(record.components[0].scores), 2)
+        self.assertEqual(record.components[0].scores[0].name, "HEPSPEC")
+        self.assertEqual(record.components[0].scores[0].factor, 1.2)
+        self.assertEqual(record.components[0].scores[1].name, "BENCHMARK")
+        self.assertEqual(record.components[0].scores[1].factor, 3.0)
+        self.assertEqual(record.components[1].name, "Memory")
+        self.assertEqual(record.components[1].amount, 100)
+        self.assertEqual(len(record.components[1].scores), 1)
+        self.assertEqual(record.components[1].scores[0].name, "BLUBB")
+        self.assertEqual(record.components[1].scores[0].factor, 1.4)