From 7a3315444c42cadaa2c634178d987a9239b40f33 Mon Sep 17 00:00:00 2001
From: Emir Muhammad <emir.muhammad@cern.ch>
Date: Wed, 25 Mar 2026 11:57:30 +0100
Subject: [PATCH] add env injector

---
 src/daqpytools/apps/logging_demonstrator.py | 91 ++++++++++++++++++---
 1 file changed, 79 insertions(+), 12 deletions(-)

diff --git a/src/daqpytools/apps/logging_demonstrator.py b/src/daqpytools/apps/logging_demonstrator.py
index 7398028..20f1640 100644
--- a/src/daqpytools/apps/logging_demonstrator.py
+++ b/src/daqpytools/apps/logging_demonstrator.py
@@ -1,6 +1,7 @@
 import logging
 import os
 import time
+from collections.abc import Mapping
 
 import click
 from rich.traceback import install as rich_traceback_install
@@ -40,6 +41,56 @@ def validate_test_configuration(
 
 
 
+class EnvInjector:
+    """Manage temporary overrides for ERS-related environment variables."""
+
+    def __init__(self, ers_envs: Mapping[str, str]) -> None:
+        """Initialize handler state and capture original environment values.
+
+        Args:
+            ers_envs (Mapping[str, str]): Environment variables and values to inject.
+        """
+        self.update_ers_envs(ers_envs)
+        self.update_original_envs()
+
+    def return_original_envs(self) -> dict[str, str | None]:
+        """Return current values for tracked environment variables.
+
+        Returns:
+            dict[str, str | None]: Mapping of variable names to their current values,
+                or None if unset.
+        """
+        return {
+            key: os.environ.get(key)
+            for key in self.ers_envs
+        }
+
+    def update_original_envs(self) -> None:
+        """Store a snapshot of current values for tracked environment variables."""
+        self.ori_envs = self.return_original_envs()
+
+    def update_ers_envs(self, ers_envs: Mapping[str, str]) -> None:
+        """Update the environment variables managed by this handler.
+
+        Args:
+            ers_envs (Mapping[str, str]): Variables and values to manage.
+        """
+        self.ers_envs = dict(ers_envs)
+
+    def inject_ers_envs(self) -> None:
+        """Inject managed environment variable values into process environment."""
+        for key, value in self.ers_envs.items():
+            os.environ[key] = value
+
+    def restore_original_envs(self) -> None:
+        """Restore tracked environment variables to their original values."""
+        for var_name, original_value in self.ori_envs.items():
+            if original_value is None:
+                os.environ.pop(var_name, None)
+            else:
+                os.environ[var_name] = original_value
+
+
 def test_main_functions(main_logger:logging.Logger) -> None:
     """Demonstrates the main functionality of the daqpytools logger.
     
@@ -193,17 +244,20 @@ def test_handlerconf(main_logger: logging.Logger) -> None:
     main_logger.warning("Handlerconf Base", extra=handlerconf.Base)
     main_logger.warning("Handlerconf Opmon", extra=handlerconf.Opmon)
 
-    #* Interlude: Inject sample environment variables
-    os.environ["DUNEDAQ_ERS_WARNING"] = "erstrace,throttle,lstdout"
-    os.environ["DUNEDAQ_ERS_INFO"] = "erstrace,throttle,lstdout"
-    os.environ["DUNEDAQ_ERS_FATAL"] = "erstrace,lstdout"
-    os.environ["DUNEDAQ_ERS_ERROR"] = (
-        "erstrace,"
-        "throttle,"
-        "lstdout,"
-        "protobufstream(monkafka.cern.ch:30092)"
-    )
-        
+    ers_envs = {
+        "DUNEDAQ_ERS_WARNING" : "erstrace,throttle,lstdout",
+        "DUNEDAQ_ERS_INFO" : "erstrace,throttle,lstdout",
+        "DUNEDAQ_ERS_FATAL" : "erstrace,lstdout",
+        "DUNEDAQ_ERS_ERROR" : (
+                "erstrace,"
+                "throttle,"
+                "lstdout,"
+                "protobufstream(monkafka.cern.ch:30092)"
+            )
+    }   
+    env_injector = EnvInjector(ers_envs)
+    env_injector.inject_ers_envs()
+    
     info_out = f"{os.getenv('DUNEDAQ_ERS_ERROR')=}"
     main_logger.info(info_out)
     critical_out = f"{os.getenv('DUNEDAQ_ERS_CRITICAL')=}"
@@ -226,6 +280,8 @@ def test_handlerconf(main_logger: logging.Logger) -> None:
         extra=handlerconf.ERS
     )     
 
+    env_injector.restore_original_envs()
+
 
 def test_fallback_handlers(log_level: str) -> None:
     """Demonstrate fallback handler behavior for a logger.
@@ -280,6 +336,15 @@ def test_ers_handler_configuration(log_level: str) -> None:
     os.environ["DUNEDAQ_ERS_FATAL"] = "lstderr,rich"
     os.environ["DUNEDAQ_ERS_ERROR"] = "rich"
 
+    ers_envs = {
+        "DUNEDAQ_ERS_WARNING" : "rich",
+        "DUNEDAQ_ERS_INFO" : "lstdout",
+        "DUNEDAQ_ERS_FATAL" : "lstderr,rich",
+        "DUNEDAQ_ERS_ERROR" : "rich"
+    }   
+    env_injector = EnvInjector(ers_envs)
+    env_injector.inject_ers_envs()
+
     ers_logger: logging.Logger = get_daq_logger(
         logger_name="ers_logger",
         log_level=log_level,
@@ -294,9 +359,11 @@ def test_ers_handler_configuration(log_level: str) -> None:
     
     ers_hc = LogHandlerConf(init_ers=True)
     ers_logger.info("ERS Info lstdout ", extra=ers_hc.ERS)
-    ers_logger.warning("ERS error lstdout", extra=ers_hc.ERS)
+    ers_logger.warning("ERS error rich", extra=ers_hc.ERS)
     ers_logger.critical("ERS critical lstderr + rich", extra=ers_hc.ERS)
 
+    env_injector.restore_original_envs()
+
 
 class AllOptionsCommand(click.Command):
     """Parse the arguments passed and validate they are acceptable, otherwise print the