diff --git a/core/Azure.Mcp.Core/src/Helpers/CommandHelper.cs b/core/Azure.Mcp.Core/src/Helpers/CommandHelper.cs index dace31fb7b..c739d5515a 100644 --- a/core/Azure.Mcp.Core/src/Helpers/CommandHelper.cs +++ b/core/Azure.Mcp.Core/src/Helpers/CommandHelper.cs @@ -26,11 +26,9 @@ public static bool HasSubscriptionAvailable(CommandResult commandResult) var subscriptionValue = parseResult.GetValueOrDefault(OptionDefinitions.Common.Subscription.Name); var envSubscription = EnvironmentHelpers.GetAzureSubscriptionId(); - return (string.IsNullOrEmpty(subscriptionValue) || IsPlaceholder(subscriptionValue)) && !string.IsNullOrEmpty(envSubscription) + return string.IsNullOrEmpty(subscriptionValue) && !string.IsNullOrEmpty(envSubscription) ? envSubscription : subscriptionValue; } - - private static bool IsPlaceholder(string value) => value.Contains("subscription") || value.Contains("default"); } } diff --git a/core/Azure.Mcp.Core/src/Services/Azure/Subscription/SubscriptionService.cs b/core/Azure.Mcp.Core/src/Services/Azure/Subscription/SubscriptionService.cs index 09a1df6f69..1660422ccb 100644 --- a/core/Azure.Mcp.Core/src/Services/Azure/Subscription/SubscriptionService.cs +++ b/core/Azure.Mcp.Core/src/Services/Azure/Subscription/SubscriptionService.cs @@ -47,8 +47,23 @@ public async Task GetSubscription(string subscription, str { ValidateRequiredParameters((nameof(subscription), subscription)); - // Get the subscription ID first, whether the input is a name or ID - var subscriptionId = await GetSubscriptionId(subscription, tenant, retryPolicy, cancellationToken); + string subscriptionId; + try + { + // Get the subscription ID first, whether the input is a name or ID + subscriptionId = await GetSubscriptionId(subscription, tenant, retryPolicy, cancellationToken); + } + catch + { + // If the subscription lookup failed, try environment variable as fallback + var envSubscription = Core.Helpers.EnvironmentHelpers.GetAzureSubscriptionId(); + if (string.IsNullOrEmpty(envSubscription)) + { + throw; // Re-throw the original exception if no environment variable is available + } + + subscriptionId = await GetSubscriptionId(envSubscription, tenant, retryPolicy, cancellationToken); + } // Use subscription ID for cache key var cacheKey = string.IsNullOrEmpty(tenant) diff --git a/core/Azure.Mcp.Core/tests/Azure.Mcp.Core.UnitTests/Services/Azure/Subscription/SubscriptionServiceTests.cs b/core/Azure.Mcp.Core/tests/Azure.Mcp.Core.UnitTests/Services/Azure/Subscription/SubscriptionServiceTests.cs new file mode 100644 index 0000000000..fbec58454c --- /dev/null +++ b/core/Azure.Mcp.Core/tests/Azure.Mcp.Core.UnitTests/Services/Azure/Subscription/SubscriptionServiceTests.cs @@ -0,0 +1,124 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using Azure.Mcp.Core.Helpers; +using Azure.Mcp.Core.Options; +using Azure.Mcp.Core.Services.Azure.Subscription; +using Azure.Mcp.Core.Services.Azure.Tenant; +using Azure.Mcp.Core.Services.Caching; +using NSubstitute; +using Xunit; + +namespace Azure.Mcp.Core.UnitTests.Services.Azure.Subscription; + +public class SubscriptionServiceTests : IDisposable +{ + private const string TestSubscriptionId = "12345678-1234-1234-1234-123456789012"; + private const string TestSubscriptionName = "Test Subscription"; + private const string EnvVarSubscriptionId = "87654321-4321-4321-4321-210987654321"; + + private readonly ICacheService _cacheService; + private readonly ITenantService _tenantService; + private readonly SubscriptionService _subscriptionService; + private readonly string? _originalEnvValue; + + public SubscriptionServiceTests() + { + // Save original environment variable + _originalEnvValue = Environment.GetEnvironmentVariable("AZURE_SUBSCRIPTION_ID"); + + _cacheService = Substitute.For(); + _tenantService = Substitute.For(); + _subscriptionService = new SubscriptionService(_cacheService, _tenantService); + + // Setup default cache behavior to return null (cache miss) + _cacheService.GetAsync(Arg.Any(), Arg.Any(), Arg.Any()) + .Returns((object?)null); + } + + public void Dispose() + { + // Restore original environment variable + if (_originalEnvValue != null) + { + Environment.SetEnvironmentVariable("AZURE_SUBSCRIPTION_ID", _originalEnvValue); + } + else + { + Environment.SetEnvironmentVariable("AZURE_SUBSCRIPTION_ID", null); + } + } + + [Fact] + public void IsSubscriptionId_ValidGuid_ReturnsTrue() + { + // Act + var result = _subscriptionService.IsSubscriptionId(TestSubscriptionId); + + // Assert + Assert.True(result); + } + + [Fact] + public void IsSubscriptionId_InvalidGuid_ReturnsFalse() + { + // Act + var result = _subscriptionService.IsSubscriptionId("not-a-guid"); + + // Assert + Assert.False(result); + } + + [Fact] + public void IsSubscriptionId_SubscriptionName_ReturnsFalse() + { + // Act + var result = _subscriptionService.IsSubscriptionId(TestSubscriptionName); + + // Assert + Assert.False(result); + } + + [Fact] + public void EnvironmentVariableFallback_IsAvailable() + { + // Arrange + Environment.SetEnvironmentVariable("AZURE_SUBSCRIPTION_ID", EnvVarSubscriptionId); + + // Act + var envSubscription = EnvironmentHelpers.GetAzureSubscriptionId(); + + // Assert + Assert.Equal(EnvVarSubscriptionId, envSubscription); + } + + [Fact] + public void EnvironmentVariableFallback_WhenNotSet_ReturnsNull() + { + // Arrange + Environment.SetEnvironmentVariable("AZURE_SUBSCRIPTION_ID", null); + + // Act + var envSubscription = EnvironmentHelpers.GetAzureSubscriptionId(); + + // Assert + Assert.Null(envSubscription); + } + + [Theory] + [InlineData("12345678-1234-1234-1234-123456789012", true)] + [InlineData("00000000-0000-0000-0000-000000000000", true)] + [InlineData("subscription", false)] + [InlineData("default", false)] + [InlineData("My Subscription Name", false)] + [InlineData("test-subscription", false)] + public void IsSubscriptionId_VariousInputs_ReturnsExpected(string input, bool expected) + { + // Act + var result = _subscriptionService.IsSubscriptionId(input); + + // Assert + Assert.Equal(expected, result); + } +} +