CommandTestsBase.cs

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using Microsoft.Mcp.Tests.Client.Helpers;
using Microsoft.Mcp.Tests.Helpers;
using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol;
using Xunit;

namespace Microsoft.Mcp.Tests.Client;

[Trait("TestType", "Live")]
public abstract class CommandTestsBase(ITestOutputHelper output, LiveServerFixture liveServerFixture)
    : IAsyncLifetime, IDisposable, IClassFixture<LiveServerFixture>
{
    protected const string TenantNameReason = "Service principals cannot use TenantName for lookup";

    protected McpClient Client { get; private set; } = default!;
    protected LiveTestSettings Settings { get; set; } = default!;
    protected StringBuilder FailureOutput { get; } = new();
    protected ITestOutputHelper Output { get; } = output;
    protected LiveServerFixture LiveServerFixture { get; } = liveServerFixture;

    public TestMode TestMode = TestMode.Live;

    /// <summary>
    /// Sets the tools that the MCP server will scope to. Call this before InitializeAsync().
    /// <para>
    /// This reduces the number of tools that the MCP server will load and initialize.
    /// </para>
    /// </summary>
    /// <param name="tools">Tools to pass to the server (e.g., "storage_blob_get", "appconfig_kv_get", etc.)</param>
    public virtual string[] Tools => [];

    public virtual async ValueTask InitializeAsync()
    {
        await InitializeAsyncInternal(null);
    }

    public static LiveTestSettings PlaybackSettings => new()
    {
        SubscriptionId = "00000000-0000-0000-0000-000000000000",
        TenantId = "00000000-0000-0000-0000-000000000000",
        ResourceBaseName = "Sanitized",
        SubscriptionName = "Sanitized",
        TenantName = "Sanitized",
        ResourceGroupName = "Sanitized",
        TestMode = TestMode.Playback
    };

    protected virtual async ValueTask LoadSettingsAsync()
    {
        Settings = await TryLoadLiveSettingsAsync().ConfigureAwait(false) ?? PlaybackSettings;

        // if the user has set to playback in LiveTestSettings, they're
        // intentionally checking playback mode, load the playback settings
        // and ignore what we got from the .testsettings.json file
        if (Settings.TestMode == TestMode.Playback)
        {
            Settings = PlaybackSettings;
        }

        TestMode = Settings.TestMode;
    }

    private static async Task<LiveTestSettings?> TryLoadLiveSettingsAsync()
    {
        try
        {
            var settingsFixture = new LiveTestSettingsFixture();
            await settingsFixture.InitializeAsync().ConfigureAwait(false);
            return settingsFixture.Settings;
        }
        catch (FileNotFoundException)
        {
            return null;
        }
    }

    private Dictionary<string, string?> GetEnvironmentVariables(TestProxyFixture? proxy)
    {
        Dictionary<string, string?> envVarDictionary = [
            // Propagate playback signaling & sanitized identifiers to server process.

            // TODO: Temporarily commenting these out until we can solve for subscription id tests
            // see https://github.com/microsoft/mcp/issues/1103
            // { "AZURE_TENANT_ID", Settings.TenantId },
            // { "AZURE_SUBSCRIPTION_ID", Settings.SubscriptionId }
        ];

        // Add any custom environment variables from settings
        if (Settings?.EnvironmentVariables != null)
        {
            foreach (var kvp in Settings.EnvironmentVariables)
            {
                envVarDictionary[kvp.Key] = kvp.Value;
            }
        }

        if (proxy != null && proxy.Proxy != null)
        {
            envVarDictionary["TEST_PROXY_URL"] = proxy.Proxy.BaseUri;

            if (TestMode is TestMode.Playback)
            {
                // AZURE_TOKEN_CREDENTIALS=PlaybackTokenCredential tells the server to use a special credential that
                // returns fake tokens in playback mode, which prevents any accidental live calls if a test is misconfigured
                envVarDictionary["AZURE_TOKEN_CREDENTIALS"] = "PlaybackTokenCredential";
                envVarDictionary["TEST_MODE"] = "Playback";
            }
        }

        return envVarDictionary;
    }

    protected virtual async ValueTask InitializeAsyncInternal(TestProxyFixture? proxy = null)
    {
        await LoadSettingsAsync();

        // Use custom arguments if provided, otherwise use standard mode (debug can be enabled via environment variable)
        var debugEnvVar = Environment.GetEnvironmentVariable("AZURE_MCP_TEST_DEBUG");
        var enableDebug = string.Equals(debugEnvVar, "true", StringComparison.OrdinalIgnoreCase) || Settings.DebugOutput;

        // Live testing uses '--mode all' as it doesn't require sampling and the way live tests are ran an agent isn't
        // used to make the request, so sampling isn't available.
        List<string> arguments = enableDebug
            ? ["server", "start", "--mode", "all", "--debug", "--dangerously-disable-elicitation", "--disable-caching"]
            : ["server", "start", "--mode", "all", "--dangerously-disable-elicitation", "--disable-caching"];

        if (Tools != null && Tools.Length > 0)
        {
            foreach (var tool in Tools)
            {
                arguments.Add("--tool");
                arguments.Add(tool);
            }
        }

        LiveServerFixture.EnvironmentVariables = GetEnvironmentVariables(proxy);
        LiveServerFixture.Arguments = arguments;
        LiveServerFixture.Output = Output;
        LiveServerFixture.Settings = Settings;

        await LiveServerFixture.EnsureStartedAsync();
        Client = LiveServerFixture.GetMcpClient();
    }

    /// <summary>
    /// Calls <see cref="McpClient.CallToolAsync(string, IReadOnlyDictionary{string, object?}?, IProgress{ModelContextProtocol.ProgressNotificationValue}?, ModelContextProtocol.RequestOptions?, CancellationToken)"/>
    /// executing the command against the MCP server and returns the "results" property from <see cref="CommandResponse"/>, if it exists.
    /// Logs the request and response for debugging purposes.
    /// </summary>
    /// <param name="command">The MCP server command to execute.</param>
    /// <param name="parameters">The MCP server command parameters.</param>
    /// <returns>The "results" JSON property from <see cref="CommandResponse"/>, if it exists.</returns>
    protected async Task<JsonElement?> CallToolAsync(string command, Dictionary<string, object?> parameters)
        => await CallToolAsync(command, parameters, Client);

    /// <summary>
    /// Calls <see cref="McpClient.CallToolAsync(string, IReadOnlyDictionary{string, object?}?, IProgress{ModelContextProtocol.ProgressNotificationValue}?, ModelContextProtocol.RequestOptions?, CancellationToken)"/>
    /// executing the command against the MCP server and returns the "results" property from <see cref="CommandResponse"/>, if it exists.
    /// Logs the request and response for debugging purposes.
    /// </summary>
    /// <param name="command">The MCP server command to execute.</param>
    /// <param name="parameters">The MCP server command parameters.</param>
    /// <param name="mcpClient">The MCP client to use for the call.</param>
    /// <returns>The "results" JSON property from <see cref="CommandResponse"/>, if it exists.</returns>
    protected async Task<JsonElement?> CallToolAsync(string command, Dictionary<string, object?> parameters, McpClient mcpClient)
        => await CallToolAsync(command, parameters, mcpClient, elem => elem.TryGetProperty("results", out var property) ? property : null);

    /// <summary>
    /// Calls <see cref="McpClient.CallToolAsync(string, IReadOnlyDictionary{string, object?}?, IProgress{ModelContextProtocol.ProgressNotificationValue}?, ModelContextProtocol.RequestOptions?, CancellationToken)"/>
    /// executing the command against the MCP server and extracts the JSON property from <see cref="CommandResponse"/>, if it exists.
    /// Logs the request and response for debugging purposes.
    /// </summary>
    /// <param name="command">The MCP server command to execute.</param>
    /// <param name="parameters">The MCP server command parameters.</param>
    /// <param name="mcpClient">The MCP client to use for the call. If null the default Client will be used.</param>
    /// <param name="resultProcessor">A function to extract the desired result from the JSON response. If null the "results" property will be retrieved, if it exists.</param>
    /// <returns>The extracted JSON property from <see cref="CommandResponse"/>, if it exists.</returns>
    protected async Task<JsonElement?> CallToolAsync(
        string command,
        Dictionary<string, object?> parameters,
        McpClient? mcpClient = null,
        Func<JsonElement, JsonElement?>? resultProcessor = null)
    {
        // Use the same debug logic as MCP server initialization
        var debugEnvVar = Environment.GetEnvironmentVariable("AZURE_MCP_TEST_DEBUG");
        var enableDebug = string.Equals(debugEnvVar, "true", StringComparison.OrdinalIgnoreCase) || Settings.DebugOutput;
        mcpClient ??= Client;
        resultProcessor ??= (elem => elem.TryGetProperty("results", out var property) ? property : null);

        // Output will be streamed, so if we're not in debug mode, hold the debug output for logging in the failure case
        Action<string> writeOutput = enableDebug
            ? Output.WriteLine
            : s => FailureOutput.AppendLine(s);

        writeOutput($"request: {JsonSerializer.Serialize(new { command, parameters })}");

        CallToolResult result;
        try
        {
            result = await mcpClient.CallToolAsync(command, parameters);
        }
        catch (ModelContextProtocol.McpException ex)
        {
            // MCP client throws exceptions for error responses, but we want to handle them gracefully
            writeOutput($"MCP exception: {ex.Message}");
            throw; // Re-throw if we can't handle it
        }

        var content = McpTestUtilities.GetFirstText(result.Content);
        if (string.IsNullOrWhiteSpace(content))
        {
            writeOutput($"response: {JsonSerializer.Serialize(result)}");
            throw new Exception("No JSON content found in the response.");
        }

        JsonElement root;
        try
        {
            root = JsonSerializer.Deserialize<JsonElement>(content!);
            if (root.ValueKind != JsonValueKind.Object)
            {
                throw new Exception("Invalid JSON response.");
            }

            // Remove the `args` property and log the content
            var trimmed = root.Deserialize<JsonObject>()!;
            trimmed.Remove("args");
            writeOutput($"response: {trimmed.ToJsonString(new JsonSerializerOptions { WriteIndented = true })}");
        }
        catch (Exception ex)
        {
            // If we can't json parse the content as a JsonObject, log the content and throw an exception
            writeOutput($"response: {content}");
            throw new Exception("Failed to deserialize JSON response.", ex);
        }

        return resultProcessor.Invoke(root);
    }

    public void Dispose()
    {
        Dispose(disposing: true);
        GC.SuppressFinalize(this);
    }

    public virtual async ValueTask DisposeAsync()
    {
        await DisposeAsyncCore().ConfigureAwait(false);
        Dispose(disposing: false);
        GC.SuppressFinalize(this);
    }

    // subclasses should override this method to dispose resources
    // subclasses should override this method to dispose InitializeAsyncresources
    // overrides should still call base.Dispose(disposing)
    protected virtual void Dispose(bool disposing)
    {
        if (disposing)
        {
            // No unmanaged resources to release, but if we had, we'd release them here.
            // _disposableResource?.Dispose();
            // _disposableResource = null;
        }

        // Failure output may contain request and response details that should be output for failed tests.
        if (TestContext.Current?.TestState?.Result == TestResult.Failed && FailureOutput.Length > 0)
        {
            Output.WriteLine(FailureOutput.ToString());
        }
    }

    // subclasses should override this method to dispose async resources
    // overrides should still call base.DisposeAsyncCore()
    protected virtual ValueTask DisposeAsyncCore() => ValueTask.CompletedTask;
}