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

@@ -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()
 
 

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(),
+}

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."""

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

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()

tests/unit/test_config.py

@@ -25,17 +25,20 @@ 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):
         """Test invalid config."""
         mock_safe_load.return_value = {
             "port": -10000,
             "level": "RANDOM",
+            "enable_collectors": ["megaraidcollector"],
         }
         with pytest.raises(ValueError):
             Config.load_config()