fix UT

Author: mainred

src/aks-agent/azext_aks_agent/custom.py

@@ -10,10 +10,7 @@
 # pylint: disable=too-many-lines, disable=broad-except
 from azext_aks_agent.agent.agent import aks_agent as aks_agent_internal
 from azext_aks_agent.agent.llm_config_manager import LLMConfigManager
-from azext_aks_agent.agent.llm_providers import (
-    PROVIDER_REGISTRY,
-    prompt_provider_choice,
-)
+from azext_aks_agent.agent.llm_providers import prompt_provider_choice
 from azext_aks_agent.agent.logging import rich_logging
 from azure.cli.core.api import get_config_dir
 from azure.cli.core.azclierror import AzCLIError

src/aks-agent/azext_aks_agent/tests/latest/test_aks_agent_llm_config_manager.py

@@ -647,7 +647,6 @@ def test_export_model_config_sets_environment_variables(self, mock_environ, mock
 
         # Verify environment variables were set
         self.assertEqual(mock_environ["OPENAI_API_KEY"], "test-key")
-        self.assertEqual("MODEL_NAME", mock_environ)
         self.assertNotIn("provider", mock_environ)
 
         # Verify provider was called correctly
@@ -685,7 +684,7 @@ def test_export_model_config_azure_converts_model_name(self, mock_environ, mock_
         self.assertEqual(llm_config["DEPLOYMENT_NAME"], "gpt-4-legacy")
 
         # Verify environment variables
-        self.assertEqual("DEPLOYMENT_NAME", mock_environ)
+        self.assertEqual(mock_environ["AZURE_OPENAI_API_KEY"], "azure-key")
         self.assertEqual(mock_environ["AZURE_OPENAI_ENDPOINT"], "https://test.openai.azure.com")
 
         # Verify provider was called with deployment name
@@ -716,7 +715,7 @@ def test_export_model_config_azure_with_deployment_name(self, mock_environ, mock
 
         # Verify DEPLOYMENT_NAME is preserved
         self.assertEqual(llm_config["DEPLOYMENT_NAME"], "my-gpt-4-deployment")
-        self.assertEqual("DEPLOYMENT_NAME", mock_environ)
+        self.assertEqual(mock_environ["AZURE_OPENAI_ENDPOINT"], "https://test.openai.azure.com")
 
         # Verify provider was called with deployment name
         mock_provider_instance.model_name.assert_called_once_with("my-gpt-4-deployment")