Fix memory leak in logging.enable

applicationinsights/logging/LoggingHandler.py

@@ -1,10 +1,11 @@
 import logging
 import applicationinsights
+from weakref import WeakValueDictionary
 from applicationinsights.channel import AsynchronousSender, AsynchronousQueue
 from applicationinsights.channel import SynchronousSender, SynchronousQueue
 from applicationinsights.channel import TelemetryChannel
 
-enabled_instrumentation_keys = {}
+enabled_instrumentation_keys = WeakValueDictionary()
 
 def enable(instrumentation_key, *args, **kwargs):
     """Enables the Application Insights logging handler for the root logger for the supplied instrumentation key.

tests/applicationinsights_tests/logging_tests/TestLoggingHandler.py

@@ -11,6 +11,7 @@
     sys.path.append(rootDirectory)
 
 from applicationinsights import logging
+from applicationinsights.logging.LoggingHandler import enabled_instrumentation_keys
 
 class TestEnable(unittest.TestCase):
     def test_enable(self):
@@ -63,6 +64,20 @@ def test_enable_with_level(self):
     def test_enable_raises_exception_on_no_instrumentation_key(self):
         self.assertRaises(Exception, logging.enable, None)
 
+    def test_handler_removal_clears_cache(self):
+        def enable_telemetry():
+            logging.enable('key1')
+
+        def remove_telemetry_handlers():
+            for handler in pylogging.getLogger().handlers:
+                if isinstance(handler, logging.LoggingHandler):
+                    pylogging.getLogger().removeHandler(handler)
+
+        enable_telemetry()
+        self.assertIn('key1', enabled_instrumentation_keys)
+        remove_telemetry_handlers()
+        self.assertNotIn('key1', enabled_instrumentation_keys)
+
 
 class TestLoggingHandler(unittest.TestCase):
     def test_construct(self):