mark_process_dead respects old env var

csmarchbanks · csmarchbanks · commit 9d745ed0f8b8 · 2021-04-08T09:56:44.000-06:00
Some users use mark_process_dead without going through any of the
__init__() logic which sets the new environment variable.

Signed-off-by: Chris Marchbanks &lt;csmarchbanks@gmail.com&gt;
diff --git a/prometheus_client/multiprocess.py b/prometheus_client/multiprocess.py
@@ -157,7 +157,7 @@ def collect(self):
 def mark_process_dead(pid, path=None):
     """Do bookkeeping for when one process dies in a multi-process setup."""
     if path is None:
-        path = os.environ.get('PROMETHEUS_MULTIPROC_DIR')
+        path = os.environ.get('PROMETHEUS_MULTIPROC_DIR', os.environ.get('prometheus_multiproc_dir'))
     for f in glob.glob(os.path.join(path, 'gauge_livesum_{0}.db'.format(pid))):
         os.remove(f)
     for f in glob.glob(os.path.join(path, 'gauge_liveall_{0}.db'.format(pid))):
diff --git a/tests/test_multiprocess.py b/tests/test_multiprocess.py
@@ -30,8 +30,8 @@ def setUp(self):
         self.tempdir = tempfile.mkdtemp()
 
     def tearDown(self):
-        del os.environ['prometheus_multiproc_dir']
-        del os.environ['PROMETHEUS_MULTIPROC_DIR']
+        os.environ.pop('prometheus_multiproc_dir', None)
+        os.environ.pop('PROMETHEUS_MULTIPROC_DIR', None)
         values.ValueClass = MutexValue
         shutil.rmtree(self.tempdir)
 
@@ -48,6 +48,12 @@ def test_deprecation_warning(self):
             assert issubclass(w[-1].category, DeprecationWarning)
             assert "PROMETHEUS_MULTIPROC_DIR" in str(w[-1].message)
 
+    def test_mark_process_dead_respects_lowercase(self):
+        os.environ['prometheus_multiproc_dir'] = self.tempdir
+        # Just test that this does not raise with a lowercase env var. The
+        # logic is tested elsewhere.
+        mark_process_dead(123)
+
 
 class TestMultiProcess(unittest.TestCase):
     def setUp(self):
