diff --git a/opnsense_exporter/server.py b/opnsense_exporter/server.py index 4339690..5bb5122 100644 --- a/opnsense_exporter/server.py +++ b/opnsense_exporter/server.py @@ -51,43 +51,52 @@ active_server_bytes_transmitted = Gauge( ) -def process_requests(main, backup, exporter_instance: str = ""): - """A dummy function that takes some time.""" - main_state = main.get_interface_vip_status() - backup_sate = backup.get_interface_vip_status() - main_ha_state.labels(instance=exporter_instance, **main.labels).state(main_state) - backup_ha_state.labels(instance=exporter_instance, **backup.labels).state( - backup_sate - ) - active_opnsense = None - if main_state == "active": - active_opnsense = main - if backup_sate == "active": - active_opnsense = backup - if active_opnsense: - bytes_received, bytes_transmitted = active_opnsense.get_wan_trafic() - if bytes_received or bytes_received == 0: - active_server_bytes_received.labels( - instance=exporter_instance, **active_opnsense.labels - ).set(bytes_received) - if bytes_transmitted or bytes_transmitted == 0: - active_server_bytes_transmitted.labels( - instance=exporter_instance, **active_opnsense.labels - ).set(bytes_transmitted) +class OPNSensePrometheusExporter: + def __init__( + self, + main: OPNSenseAPI, + backup: OPNSenseAPI, + exporter_instance: str = "", + check_frequency: int = 1, + ): + self.main = main + self.backup = backup + self.exporter_instance = exporter_instance + self.check_frequency = check_frequency + def process_requests(self): + """A dummy function that takes some time.""" + main_state = self.main.get_interface_vip_status() + backup_sate = self.backup.get_interface_vip_status() + main_ha_state.labels(instance=self.exporter_instance, **self.main.labels).state( + main_state + ) + backup_ha_state.labels( + instance=self.exporter_instance, **self.backup.labels + ).state(backup_sate) + active_opnsense = None + if main_state == "active": + active_opnsense = self.main + if backup_sate == "active": + active_opnsense = self.backup + if active_opnsense: + bytes_received, bytes_transmitted = active_opnsense.get_wan_trafic() + if bytes_received or bytes_received == 0: + active_server_bytes_received.labels( + instance=self.exporter_instance, **active_opnsense.labels + ).set(bytes_received) + if bytes_transmitted or bytes_transmitted == 0: + active_server_bytes_transmitted.labels( + instance=self.exporter_instance, **active_opnsense.labels + ).set(bytes_transmitted) -def start_server( - main: OPNSenseAPI, - backup: OPNSenseAPI, - check_frequency: int = 1, - exporter_instance: str = "", -): - # Start up the server to expose the metrics. - start_http_server(8000) - # Generate some requests. - while True: - process_requests(main, backup, exporter_instance=exporter_instance) - time.sleep(check_frequency) + def start_server(self): + # Start up the server to expose the metrics. + start_http_server(8000) + # Generate some requests. + while True: + self.process_requests() + time.sleep(self.check_frequency) def run(): @@ -147,7 +156,8 @@ def run(): ) arguments = parser.parse_args() - start_server( + + server = OPNSensePrometheusExporter( OPNSenseAPI( OPNSenseRole.MAIN, arguments.main, arguments.user, arguments.password ), @@ -157,3 +167,7 @@ def run(): check_frequency=arguments.frequency, exporter_instance=arguments.prom_instance, ) + server.start_server() + + # return the server instance mainly for test purpose + return server diff --git a/tests/test_server.py b/tests/test_server.py index 2774776..3219055 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -3,7 +3,7 @@ from unittest import mock import responses from opnsense_exporter.opnsense_api import OPNSenseAPI, OPNSenseRole -from opnsense_exporter.server import process_requests, run +from opnsense_exporter.server import OPNSensePrometheusExporter, run from .common import ( BACKUP_HOST, @@ -41,7 +41,7 @@ class FakePromGauge(FakePromMetric): self.value = value -@mock.patch("opnsense_exporter.server.start_server") +@mock.patch("opnsense_exporter.server.OPNSensePrometheusExporter.start_server") def test_parser(server_mock): with mock.patch( "sys.argv", @@ -61,16 +61,18 @@ def test_parser(server_mock): "server-hostname-instance", ], ): - run() + server = run() server_mock.assert_called_once() - main, bck = server_mock.call_args.args - assert main.login == "user-test" - assert bck.login == "user-test" - assert main.password == "pwd-test" - assert bck.password == "pwd-test" - assert main.host == "main.host" - assert bck.host == "backup.host" - assert server_mock.call_args.kwargs["check_frequency"] == 15 + + assert server.main.role == OPNSenseRole.MAIN + assert server.main.host == "main.host" + assert server.main.login == "user-test" + assert server.main.password == "pwd-test" + assert server.backup.role == OPNSenseRole.BACKUP + assert server.backup.host == "backup.host" + assert server.backup.login == "user-test" + assert server.backup.password == "pwd-test" + assert server.check_frequency == 15 @responses.activate @@ -108,10 +110,10 @@ def test_process_requests(): "opnsense_exporter.server.active_server_bytes_transmitted", new=active_server_bytes_transmitted_mock, ): - process_requests( + OPNSensePrometheusExporter( OPNSenseAPI(OPNSenseRole.MAIN, MAIN_HOST, LOGIN, PASSWORD), OPNSenseAPI(OPNSenseRole.BACKUP, BACKUP_HOST, LOGIN, PASSWORD), - ) + ).process_requests() assert main_ha_state_mock._state == "active" assert main_ha_state_mock.count_state_calls == 1 @@ -181,10 +183,10 @@ def test_process_requests_backup_active(): "opnsense_exporter.server.active_server_bytes_transmitted", new=active_server_bytes_transmitted_mock, ): - process_requests( + OPNSensePrometheusExporter( OPNSenseAPI(OPNSenseRole.MAIN, MAIN_HOST, LOGIN, PASSWORD), OPNSenseAPI(OPNSenseRole.BACKUP, BACKUP_HOST, LOGIN, PASSWORD), - ) + ).process_requests() assert main_ha_state_mock._state == "maintenancemode" assert main_ha_state_mock.count_state_calls == 1 assert main_ha_state_mock._labels == { @@ -254,10 +256,10 @@ def test_process_no_active(): "opnsense_exporter.server.active_server_bytes_transmitted", new=active_server_bytes_transmitted_mock, ): - process_requests( + OPNSensePrometheusExporter( OPNSenseAPI(OPNSenseRole.MAIN, MAIN_HOST, LOGIN, PASSWORD), OPNSenseAPI(OPNSenseRole.BACKUP, BACKUP_HOST, LOGIN, PASSWORD), - ) + ).process_requests() assert main_ha_state_mock._state == "maintenancemode" assert main_ha_state_mock.count_state_calls == 1 @@ -315,10 +317,10 @@ def test_process_with_falsy_value(): "opnsense_exporter.server.active_server_bytes_transmitted", new=active_server_bytes_transmitted_mock, ): - process_requests( + OPNSensePrometheusExporter( OPNSenseAPI(OPNSenseRole.MAIN, MAIN_HOST, LOGIN, PASSWORD), OPNSenseAPI(OPNSenseRole.BACKUP, BACKUP_HOST, LOGIN, PASSWORD), - ) + ).process_requests() assert main_ha_state_mock._state == "active" assert main_ha_state_mock.count_state_calls == 1 assert main_ha_state_mock._labels == {