fix: improve logging

main.py

@@ -5,20 +5,22 @@
 import logging
 import os
 import sys
+from logging import Logger
 from typing import Dict, Any
 
 from config.config import get_configs
 from prometheus_client import start_http_server
 from src.exporter import MetricExporter
 
 
-def main(config: Dict[str, Any]) -> None:
+def main(config: Dict[str, Any], logger: Logger) -> None:
     app_metrics = MetricExporter(
         polling_interval_seconds=config["polling_interval_seconds"],
         dd_api_key=config["dd_api_key"],
         dd_app_key=config["dd_app_key"],
         dd_host=config["dd_host"],
         dd_debug=config["dd_debug"],
+        logger=logger,
     )
     start_http_server(config["exporter_port"])
     app_metrics.run_metrics_loop()
@@ -28,4 +30,5 @@ def main(config: Dict[str, Any]) -> None:
     config = get_configs()
     logger_format = "%(asctime)-15s %(levelname)-8s %(message)s"
     logging.basicConfig(level=config["log_level"], format=logger_format)
-    main(config)
+    logger = logging.getLogger("datadog-cost-exporter")
+    main(config, logger)

src/exporter.py

@@ -3,7 +3,6 @@
 # -*- coding:utf-8 -*-
 # Filename: exporter.py
 
-import logging
 import time
 from datetime import datetime
 from dateutil.relativedelta import relativedelta
@@ -29,8 +28,6 @@
     Info,
 )
 
-logger = logging.getLogger("datadog-cost-exporter")
-
 
 class MetricExporter:
     """
@@ -39,7 +36,13 @@ class MetricExporter:
     """
 
     def __init__(
-        self, polling_interval_seconds, dd_api_key, dd_app_key, dd_host, dd_debug
+        self,
+        polling_interval_seconds,
+        dd_api_key,
+        dd_app_key,
+        dd_host,
+        dd_debug,
+        logger,
     ):
         self.polling_interval_seconds = polling_interval_seconds
         self.dd_api_key = dd_api_key
@@ -48,6 +51,7 @@ def __init__(
         self.dd_debug = dd_debug
         self.default_labels = {}
         self.metrics = {}
+        self.logger = logger
 
     def create_api_client(self) -> ApiClient:
         configuration = Configuration()
@@ -98,14 +102,14 @@ def add_metrics(
             if metric_name not in self.metrics:
                 self.metrics[metric_name] = metric_object
             else:
-                logging.warning(f"Metric '{metric_name}' already added.")
+                self.logger.warning(f"Metric '{metric_name}' already added.")
 
     def run_metrics_loop(self):
         while True:
             try:
                 self.fetch()
             except Exception as e:
-                logging.error(e)
+                self.logger.error(e)
                 continue
             time.sleep(self.polling_interval_seconds)
 
@@ -169,7 +173,7 @@ def get_projected_total_cost(
                     }
                     metric_object.labels(**labels).set(charge.cost)
         except Exception as e:
-            logging.error(f"Error querying Datadog projected cost endpoint: {e}")
+            self.logger.error(f"Error querying Datadog projected cost endpoint: {e}")
             return None
 
     def get_historical_cost_by_org(self, api_instance: UsageMeteringApi) -> None:
@@ -213,7 +217,7 @@ def get_historical_cost_by_org(self, api_instance: UsageMeteringApi) -> None:
 
                     metric_object.labels(**labels).set(charge_cost)
         except Exception as e:
-            logging.error(f"Error handling historical cost by organization: {e}")
+            self.logger.error(f"Error handling historical cost by organization: {e}")
             return None
 
     def get_monthly_cost_attribution(self, api_instance: UsageMeteringApi) -> None:
@@ -279,7 +283,7 @@ def get_monthly_cost_attribution(self, api_instance: UsageMeteringApi) -> None:
                 metric_object.labels(**agg_type).set(value)
 
         except Exception as e:
-            logging.error(f"Error handling monthly cost attribution: {e}")
+            self.logger.error(f"Error handling monthly cost attribution: {e}")
             return None
 
     def _query_timeseries_points(
@@ -296,7 +300,7 @@ def _query_timeseries_points(
                 query=query,
             )
         except Exception as e:
-            logging.error(f"Error handling metrics list: {e}")
+            self.logger.error(f"Error handling metrics list: {e}")
             return None
 
     def get_usage_this_month(self, api_instance: V1MetricsApi) -> None:
@@ -358,7 +362,7 @@ def fetch(self):
         """
         Fetch metrics from the Datadog API.
         """
-        logger.info("Collecting metrics for a Prometheus client")
+        self.logger.info("Collecting metrics for a Prometheus client")
 
         usage_metric_api_instance = self.get_usage_metric_api_instance()
         self.get_projected_total_cost(usage_metric_api_instance)

tests/test_datadog_cost_exporter.py

@@ -16,12 +16,14 @@ def setUp(self):
         self.dd_host = "your_dd_host"
         self.dd_debug = True
         self.polling_interval_seconds = 60
+        self.logger = MagicMock()
         self.metric_exporter = MetricExporter(
             self.polling_interval_seconds,
             self.dd_api_key,
             self.dd_app_key,
             self.dd_host,
             self.dd_debug,
+            self.logger,
         )
 
     def test_create_api_client(self):
@@ -96,6 +98,7 @@ def test_fetch(self, mock_gauge, mock_usage_metering_api):
             dd_app_key="fake_app_key",
             dd_host="fake_host",
             dd_debug=True,
+            logger=MagicMock(),
         )
 
         api_instance.get_projected_cost.return_value = MagicMock()