Update tests

sdk/ml/azure-ai-ml/azure/ai/ml/_ml_client.py

@@ -10,7 +10,7 @@
 from functools import singledispatch
 from itertools import product
 from pathlib import Path
-from typing import Any, Optional, Tuple, TypeVar, Union
+from typing import Any, Dict, Optional, Tuple, TypeVar, Union
 
 from azure.ai.ml._azure_environments import (
     CloudArgumentKeys,
@@ -274,7 +274,7 @@ def __init__(
             **{"properties": properties},
             enable_telemetry=self._operation_config.enable_telemetry,
         )
-        app_insights_handler_kwargs = {"app_insights_handler": app_insights_handler}
+        app_insights_handler_kwargs: Dict[str, Tuple] = {"app_insights_handler": app_insights_handler}
 
         base_url = _get_base_url_from_metadata(cloud_name=cloud_name, is_local_mfe=True)
         self._base_url = base_url

sdk/ml/azure-ai-ml/azure/ai/ml/_telemetry/logging_handler.py

@@ -9,6 +9,7 @@
 import logging
 import platform
 import traceback
+from typing import Optional, Tuple, Union
 
 from opencensus.ext.azure.common import utils
 from opencensus.ext.azure.common.protocol import Data, Envelope, ExceptionData, Message
@@ -90,7 +91,7 @@ def get_appinsights_log_handler(
     component_name=None,
     enable_telemetry=True,
     **kwargs,
-):
+) -> Tuple[Union["AzureMLSDKLogHandler", logging.NullHandler], Optional[Tracer]]:
     """Enable the OpenCensus logging handler for specified logger and instrumentation key to send info to AppInsights.
 
     :param user_agent: Information about the user's browser.
@@ -106,21 +107,21 @@ def get_appinsights_log_handler(
     :keyword kwargs: Optional keyword arguments for adding additional information to messages.
     :paramtype kwargs: dict
     :return: The logging handler and tracer.
-    :rtype: Tuple[AzureMLSDKLogHandler, opencensus.trace.tracer.Tracer]
+    :rtype: Tuple[Union[AzureMLSDKLogHandler, logging.NullHandler], Optional[opencensus.trace.tracer.Tracer]]
     """
     try:
         if instrumentation_key is None:
             instrumentation_key = INSTRUMENTATION_KEY
 
         if not in_jupyter_notebook() or not enable_telemetry:
-            return logging.NullHandler()
+            return (logging.NullHandler(), None)
 
         if not user_agent or not user_agent.lower() == USER_AGENT.lower():
-            return logging.NullHandler()
+            return (logging.NullHandler(), None)
 
         if "properties" in kwargs and "subscription_id" in kwargs.get("properties"):
             if kwargs.get("properties")["subscription_id"] in test_subscriptions:
-                return logging.NullHandler()
+                return (logging.NullHandler(), None)
 
         child_namespace = component_name or __name__
         current_logger = logging.getLogger(AML_INTERNAL_LOGGER_NAMESPACE).getChild(child_namespace)
@@ -143,10 +144,10 @@ def get_appinsights_log_handler(
             sampler=ProbabilitySampler(1.0),
         )
 
-        return handler, tracer
+        return (handler, tracer)
     except Exception:  # pylint: disable=broad-except
         # ignore any exceptions, telemetry collection errors shouldn't block an operation
-        return logging.NullHandler(), None
+        return (logging.NullHandler(), None)
 
 
 # cspell:ignore AzureMLSDKLogHandler

sdk/ml/azure-ai-ml/azure/ai/ml/_utils/_logger_utils.py

@@ -29,8 +29,6 @@ def __init__(self, name: str):
 
     def update_info(self, data: Dict) -> None:
         if "app_insights_handler" in data:
-            app_insights_handler = data.pop("app_insights_handler")
-            if not isinstance(app_insights_handler, logging.NullHandler):
-                logger, tracer = app_insights_handler
-                self.package_logger.addHandler(logger)
-                self.package_tracer = tracer
+            logger, tracer = data.pop("app_insights_handler")
+            self.package_logger.addHandler(logger)
+            self.package_tracer = tracer

sdk/ml/azure-ai-ml/tests/internal_utils/unittests/test_logger_utils.py

@@ -3,6 +3,7 @@
 import pytest
 from mock import patch
 from opencensus.ext.azure.log_exporter import AzureLogHandler
+from opencensus.trace.tracer import Tracer
 
 from azure.ai.ml._telemetry import AML_INTERNAL_LOGGER_NAMESPACE, get_appinsights_log_handler
 from azure.ai.ml._telemetry.logging_handler import AzureMLSDKLogHandler
@@ -29,13 +30,15 @@ def test_initialize_logger_info(self) -> None:
 class TestLoggingHandler:
     def test_logging_enabled(self) -> None:
         with patch("azure.ai.ml._telemetry.logging_handler.in_jupyter_notebook", return_value=False):
-            handler, _ = get_appinsights_log_handler(user_agent=USER_AGENT)
+            handler, tracer = get_appinsights_log_handler(user_agent=USER_AGENT)
             assert isinstance(handler, logging.NullHandler)
+            assert tracer is None
 
         with patch("azure.ai.ml._telemetry.logging_handler.in_jupyter_notebook", return_value=True):
-            handler, _ = get_appinsights_log_handler(user_agent=USER_AGENT)
+            handler, tracer = get_appinsights_log_handler(user_agent=USER_AGENT)
             assert isinstance(handler, AzureLogHandler)
             assert isinstance(handler, AzureMLSDKLogHandler)
+            assert isinstance(tracer, Tracer)
 
 
 @pytest.mark.unittest
@@ -52,12 +55,12 @@ def test_init(self) -> None:
 
     def test_update_info(self) -> None:
         test_name = "test"
-        test_handler = logging.NullHandler()
+        test_handler = (logging.NullHandler(), None)
         test_data = {"app_insights_handler": test_handler}
 
         test_logger = OpsLogger(name=test_name)
         test_logger.update_info(test_data)
 
         assert len(test_data) == 0
         assert test_logger.package_logger.hasHandlers()
-        assert test_logger.package_logger.handlers[0] == test_handler
+        assert test_logger.package_logger.handlers[0] == test_handler[0]