Add diagnostics for AKS (#40988)

* Enable diagnostic logging for AKS. Needs tests

* stash

* Tests pass

* changelog

* spelling

* test_diagnostics_aks_attach not passing

* Default return

* test_diagnostics_aks_attach failing. Adding asserts

* change tests

* add more test variety

* mock exporter

* add mock args

* lint

* clarify comment

* cover functions appsvc overlap

* typo

Author: jeremydvoss

sdk/monitor/azure-monitor-opentelemetry/CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Features Added
 
+- Enabled diagnostics for AKS Attach
+    ([#40988](https://github.com/Azure/azure-sdk-for-python/pull/40988))
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/monitor/azure-monitor-opentelemetry/azure/monitor/opentelemetry/_diagnostics/diagnostic_logging.py

@@ -11,10 +11,10 @@
 
 from azure.monitor.opentelemetry._utils import (
     _EXTENSION_VERSION,
-    _IS_DIAGNOSTICS_ENABLED,
     _env_var_or_default,
     _get_customer_ikey_from_env_var,
     _get_log_path,
+    _is_diagnostics_enabled,
 )
 from azure.monitor.opentelemetry._version import VERSION
 
@@ -51,7 +51,7 @@ class AzureDiagnosticLogging:
     def _initialize(cls):
         with AzureDiagnosticLogging._lock:
             if not AzureDiagnosticLogging._initialized:
-                if _IS_DIAGNOSTICS_ENABLED and _DIAGNOSTIC_LOG_PATH:
+                if _is_diagnostics_enabled() and _DIAGNOSTIC_LOG_PATH:
                     log_format = (
                         "{"
                         + '"time":"%(asctime)s.%(msecs)03d", '

sdk/monitor/azure-monitor-opentelemetry/azure/monitor/opentelemetry/_diagnostics/status_logger.py

@@ -11,9 +11,9 @@
 
 from azure.monitor.opentelemetry._utils import (
     _EXTENSION_VERSION,
-    _IS_DIAGNOSTICS_ENABLED,
     _get_customer_ikey_from_env_var,
     _get_log_path,
+    _is_diagnostics_enabled,
 )
 from azure.monitor.opentelemetry._version import VERSION
 
@@ -45,7 +45,7 @@ def _get_status_json(cls, agent_initialized_successfully, pid, reason=None, sdk_
 
     @classmethod
     def log_status(cls, agent_initialized_successfully, reason=None, sdk_present=None):
-        if _IS_DIAGNOSTICS_ENABLED and _STATUS_LOG_PATH:
+        if _is_diagnostics_enabled() and _STATUS_LOG_PATH:
             pid = getpid()
             status_json = AzureStatusLogger._get_status_json(agent_initialized_successfully, pid, reason, sdk_present)
             if not exists(_STATUS_LOG_PATH):

sdk/monitor/azure-monitor-opentelemetry/azure/monitor/opentelemetry/_utils/__init__.py

@@ -14,6 +14,9 @@
 )
 from azure.monitor.opentelemetry.exporter._utils import (  # pylint: disable=import-error,no-name-in-module
     _is_on_app_service,
+    _is_on_aks,
+    _is_on_functions,
+    _is_attach_enabled,
 )
 from azure.monitor.opentelemetry._constants import (
     _LOG_PATH_LINUX,
@@ -26,8 +29,6 @@
 
 # --------------------Diagnostic/status logging------------------------------
 
-# TODO: Add environment variable to enabled diagnostics off of App Service
-_IS_DIAGNOSTICS_ENABLED = _is_on_app_service()
 _CUSTOMER_IKEY_ENV_VAR = None
 
 
@@ -43,6 +44,15 @@ def _get_customer_ikey_from_env_var():
     return _CUSTOMER_IKEY_ENV_VAR
 
 
+# TODO: Add environment variable to enable/disable diagnostics
+def _is_diagnostics_enabled():
+    if _is_on_functions():
+        return False
+    if _is_on_app_service() or _is_on_aks():
+        return _is_attach_enabled()
+    return False
+
+
 def _get_log_path(status_log_path=False):
     system = platform.system()
     if system == "Linux":

sdk/monitor/azure-monitor-opentelemetry/tests/diagnostics/test_diagnostic_logging.py

@@ -92,8 +92,8 @@ def set_up(
         TEST_VERSION,
     ).start()
     patch(
-        "azure.monitor.opentelemetry._diagnostics.diagnostic_logging._IS_DIAGNOSTICS_ENABLED",
-        is_diagnostics_enabled,
+        "azure.monitor.opentelemetry._diagnostics.diagnostic_logging._is_diagnostics_enabled",
+        return_value=is_diagnostics_enabled,
     ).start()
 
 

sdk/monitor/azure-monitor-opentelemetry/tests/diagnostics/test_status_logger.py

@@ -50,8 +50,8 @@ def set_up(file_path, is_diagnostics_enabled=True):
         TEST_VERSION,
     ).start()
     patch(
-        "azure.monitor.opentelemetry._diagnostics.status_logger._IS_DIAGNOSTICS_ENABLED",
-        is_diagnostics_enabled,
+        "azure.monitor.opentelemetry._diagnostics.status_logger._is_diagnostics_enabled",
+        return_value=is_diagnostics_enabled,
     ).start()
     patch(
         "azure.monitor.opentelemetry._diagnostics.status_logger.getpid",

sdk/monitor/azure-monitor-opentelemetry/tests/utils/test_utils.py

@@ -10,6 +10,7 @@
 from unittest.mock import patch
 
 from azure.monitor.opentelemetry import _utils
+from azure.monitor.opentelemetry.exporter._constants import _AKS_ARM_NAMESPACE_ID
 
 TEST_VALUE = "TEST_VALUE"
 TEST_IKEY = "1234abcd-ab12-34cd-ab12-a23456abcdef"
@@ -45,15 +46,82 @@ def test_ikey_defaults(self):
         reload(_utils)
         self.assertEqual(_utils._get_customer_ikey_from_env_var(), "unknown")
 
-    @patch.dict("os.environ", {"WEBSITE_SITE_NAME": TEST_VALUE})
-    def test_diagnostics_enabled(self):
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_attach_enabled",
+        return_value=True,
+    )
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_on_app_service",
+        return_value=True,
+    )
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_on_aks",
+        return_value=False,
+    )
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_on_functions",
+        return_value=False,
+    )
+    def test_diagnostics_app_service_attach(
+        self, attach_mock, app_service_mock, aks_mock, functions_mock
+    ):
+        reload(_utils)
+        self.assertTrue(_utils._is_diagnostics_enabled())
+
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_attach_enabled",
+        return_value=True,
+    )
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_on_app_service",
+        return_value=False,
+    )
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_on_aks",
+        return_value=True,
+    )
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_on_functions",
+        return_value=False,
+    )
+    def test_diagnostics_aks_attach(
+        self, attach_mock, app_service_mock, aks_mock, functions_mock
+    ):
         reload(_utils)
-        self.assertTrue(_utils._IS_DIAGNOSTICS_ENABLED)
+        self.assertTrue(_utils._is_diagnostics_enabled())
 
-    def test_diagnostics_disabled(self):
-        clear_env_var("WEBSITE_SITE_NAME")
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_attach_enabled",
+        return_value=True,
+    )
+    # Functions have the WEBSITE_SITE_NAME environment variable.
+    # This causes them to appear as App Service resources
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_on_app_service",
+        return_value=True,
+    )
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_on_aks",
+        return_value=False,
+    )
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_on_functions",
+        return_value=True,
+    )
+    def test_diagnostics_functions_attach(
+        self, attach_mock, app_service_mock, aks_mock, functions_mock
+    ):
+        reload(_utils)
+        # Functions attach does not currently enable diagnostics
+        self.assertFalse(_utils._is_diagnostics_enabled())
+
+    @patch(
+        "azure.monitor.opentelemetry.exporter._utils._is_attach_enabled",
+        return_value=False,
+    )
+    def test_diagnostics_disabled(self, attach_mock):
         reload(_utils)
-        self.assertFalse(_utils._IS_DIAGNOSTICS_ENABLED)
+        self.assertFalse(_utils._is_diagnostics_enabled())
 
     @patch(
         "azure.monitor.opentelemetry._utils.platform.system",