diff --git a/tardis/plugins/auditor.py b/tardis/plugins/auditor.py index dc5346b5..525d6ec1 100644 --- a/tardis/plugins/auditor.py +++ b/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) diff --git a/tests/plugins_t/test_auditor.py b/tests/plugins_t/test_auditor.py new file mode 100644 index 00000000..e2265d07 --- /dev/null +++ b/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)