From 860e5b3a2698aa75d2479028770076f9552fbc84 Mon Sep 17 00:00:00 2001 From: Ray Chan <72372217+chanchiwai-ray@users.noreply.github.com> Date: Wed, 28 Jun 2023 09:27:25 +0800 Subject: [PATCH] Allow enabling collectors via configuration option. (#11) * Allow enabling collectors via configuration option. * Add pflake8 ignore for all pydantic.validator classmethod; and disable pylint E0213 for config.py. * Update docstring in config.py --- prometheus_hardware_exporter/__main__.py | 27 ++++++------------ prometheus_hardware_exporter/collector.py | 14 ++++++++++ prometheus_hardware_exporter/config.py | 34 ++++++++++++++++++++--- pyproject.toml | 1 + tests/unit/test_cli.py | 34 +++++++++++++++++++++++ tests/unit/test_config.py | 7 +++-- 6 files changed, 93 insertions(+), 24 deletions(-) diff --git a/prometheus_hardware_exporter/__main__.py b/prometheus_hardware_exporter/__main__.py index 2641517..149e3b0 100644 --- a/prometheus_hardware_exporter/__main__.py +++ b/prometheus_hardware_exporter/__main__.py @@ -3,19 +3,11 @@ import argparse import logging -from .collector import ( - IpmiDcmiCollector, - IpmiSelCollector, - IpmiSensorsCollector, - LSISASControllerCollector, - MegaRAIDCollector, - PowerEdgeRAIDCollector, - SsaCLICollector, -) +from .collector import COLLECTOR_REGISTRIES from .config import DEFAULT_CONFIG, Config from .exporter import Exporter -root_logger = logging.getLogger() +logger = logging.getLogger(__name__) def parse_command_line() -> argparse.Namespace: @@ -40,17 +32,16 @@ def main() -> None: """Start the prometheus-hardware-exporter.""" args = parse_command_line() config = Config.load_config(config_file=args.config or DEFAULT_CONFIG) + + root_logger = logging.getLogger() root_logger.setLevel(logging.getLevelName(config.level)) exporter = Exporter(config.port) - exporter.register(MegaRAIDCollector()) - exporter.register(IpmiDcmiCollector()) - exporter.register(IpmiSensorsCollector()) - exporter.register(IpmiSelCollector()) - exporter.register(LSISASControllerCollector(2)) - exporter.register(LSISASControllerCollector(3)) - exporter.register(PowerEdgeRAIDCollector()) - exporter.register(SsaCLICollector()) + enabled_collectors = set(config.enable_collectors) + for name, collector in COLLECTOR_REGISTRIES.items(): + if name.lower() in enabled_collectors: + logger.info("collector %s is enabled", name) + exporter.register(collector) exporter.run() diff --git a/prometheus_hardware_exporter/collector.py b/prometheus_hardware_exporter/collector.py index 0d7b7e3..154bb50 100644 --- a/prometheus_hardware_exporter/collector.py +++ b/prometheus_hardware_exporter/collector.py @@ -16,6 +16,8 @@ logger = getLogger(__name__) +__all__ = ["COLLECTOR_REGISTRIES"] + class PowerEdgeRAIDCollector(BlockingCollector): """Collector for PowerEdge RAID controller.""" @@ -832,3 +834,15 @@ def fetch(self) -> List[Payload]: def process(self, payloads: List[Payload], datastore: Dict[str, Payload]) -> List[Payload]: """Process the payload if needed.""" return payloads + + +COLLECTOR_REGISTRIES = { + "mega-raid-collector": MegaRAIDCollector(), + "ipmi-dcmi-collector": IpmiDcmiCollector(), + "ipmi-sensor-collector": IpmiSensorsCollector(), + "ipmi-sel-collector": IpmiSelCollector(), + "lsi-sas-2-collector": LSISASControllerCollector(2), + "lsi-sas-3-collector": LSISASControllerCollector(3), + "poweredge-raid-collector": PowerEdgeRAIDCollector(), + "hpe-ssa-collector": SsaCLICollector(), +} diff --git a/prometheus_hardware_exporter/config.py b/prometheus_hardware_exporter/config.py index c1b5f6b..28ad0bf 100644 --- a/prometheus_hardware_exporter/config.py +++ b/prometheus_hardware_exporter/config.py @@ -1,7 +1,8 @@ -"""Module for hardware related configuration.""" +"""Module for hardware exporter related configuration.""" import os from logging import getLogger +from typing import List from pydantic import BaseModel, validator from yaml import safe_load @@ -11,14 +12,18 @@ DEFAULT_CONFIG = os.path.join(os.environ.get("SNAP_DATA", "./"), "config.yaml") +# pylint: disable=E0213 + + class Config(BaseModel): - """Juju backup all configuration.""" + """Hardware exporter configuration.""" port: int = 10000 level: str = "DEBUG" + enable_collectors: List[str] = [] @validator("port") - def validate_port_range(cls, port: int) -> int: # noqa: N805 pylint: disable=E0213 + def validate_port_range(cls, port: int) -> int: """Validate port range.""" if not 1 <= port <= 65535: msg = "Port must be in [1, 65535]." @@ -27,7 +32,7 @@ def validate_port_range(cls, port: int) -> int: # noqa: N805 pylint: disable=E0 return port @validator("level") - def validate_level_choice(cls, level: str) -> str: # noqa: N805 pylint: disable=E0213 + def validate_level_choice(cls, level: str) -> str: """Validate logging level choice.""" level = level.upper() choices = {"DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"} @@ -37,6 +42,27 @@ def validate_level_choice(cls, level: str) -> str: # noqa: N805 pylint: disable raise ValueError(msg) return level + @validator("enable_collectors") + def validate_enable_collector_choice(cls, enable_collectors: List[str]) -> List[str]: + """Validate enable choice.""" + choices = { + "mega-raid-collector", + "ipmi-dcmi-collector", + "ipmi-sensor-collector", + "ipmi-sel-collector", + "lsi-sas-2-collector", + "lsi-sas-3-collector", + "poweredge-raid-collector", + "hpe-ssa-collector", + } + collectors = {collector.lower() for collector in enable_collectors} + invalid_choices = collectors.difference(choices) + if invalid_choices: + msg = f"{collectors} must be in {choices} (case-insensitive)." + logger.error(msg) + raise ValueError(msg) + return enable_collectors + @classmethod def load_config(cls, config_file: str = DEFAULT_CONFIG) -> "Config": """Load configuration file and validate it.""" diff --git a/pyproject.toml b/pyproject.toml index 943db6c..482ee3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,6 +7,7 @@ ignore = ["C901", "D100", "D101", "D102", "D103", "W503", "W504"] exclude = ['.eggs', '.git', '.tox', '.venv', '.build', 'build', 'report'] max-line-length = 99 max-complexity = 10 +classmethod-decorators = ["classmethod", "validator"] [tool.black] line-length = 99 diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index d94daab..9c44b4d 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -1,7 +1,10 @@ from unittest.mock import Mock, patch +import pytest + from prometheus_hardware_exporter import __main__ from prometheus_hardware_exporter.__main__ import main, parse_command_line +from prometheus_hardware_exporter.config import Config class TestCli: @@ -30,3 +33,34 @@ def test_cli_main( mock_main_parse_command_line.assert_called_once() mock_config.load_config.assert_called_once() mock_exporter.assert_called_once() + + @patch.object(__main__, "parse_command_line") + @patch.object(__main__, "Exporter") + @patch.object(__main__, "Config") + def test_cli_main_with_config(self, mock_config, mock_exporter, mock_main_parse_command_line): + mock_config.load_config.return_value = Config( + port=10000, + level="INFO", + enable_collectors=["mega-raid-collector"], + ) + mock_main_parse_command_line.return_value = Mock() + main() + mock_main_parse_command_line.assert_called_once() + mock_exporter.assert_called_once() + + @patch.object(__main__, "parse_command_line") + @patch.object(__main__, "Exporter") + @patch.object(__main__, "Config") + def test_cli_main_with_wrong_config( + self, mock_config, mock_exporter, mock_main_parse_command_line + ): + with pytest.raises(ValueError): + mock_main_parse_command_line.return_value = Mock() + mock_config.load_config.return_value = Config( + port=10000, + level="INFO", + enable_collectors=["megaraidcollector"], + ) + main() + mock_main_parse_command_line.assert_called_once() + mock_exporter.assert_called_once() diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 6c4784b..742aeeb 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -25,10 +25,12 @@ def test_valid_config(self, mock_safe_load): mock_safe_load.return_value = { "port": 10000, "level": "INFO", + "enable_collectors": ["mega-raid-collector"], } config = Config.load_config() - assert config.port == 10000 - assert config.level == "INFO" + self.assertEqual(config.port, 10000) + self.assertEqual(config.level, "INFO") + self.assertEqual(config.enable_collectors, ["mega-raid-collector"]) @patch("prometheus_hardware_exporter.config.safe_load") def test_invalid_config(self, mock_safe_load): @@ -36,6 +38,7 @@ def test_invalid_config(self, mock_safe_load): mock_safe_load.return_value = { "port": -10000, "level": "RANDOM", + "enable_collectors": ["megaraidcollector"], } with pytest.raises(ValueError): Config.load_config()