ConnectionTokenTests.cs

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *--------------------------------------------------------------------------------------------*/

using GitHub.Copilot.Test.Harness;
using Xunit;

namespace GitHub.Copilot.Test;

/// <summary>
/// Custom fixture that spawns a CLI in TCP mode with an explicit connection token, so
/// sibling clients can attempt to connect to the same port with the right/wrong/no token.
/// </summary>
public class ConnectionTokenTestFixture : IAsyncLifetime
{
    public E2ETestContext Ctx { get; private set; } = null!;
    public CopilotClient GoodClient { get; private set; } = null!;
    public int Port { get; private set; }

    public const string Token = "right-token";

    public async Task InitializeAsync()
    {
        Ctx = await E2ETestContext.CreateAsync();
        GoodClient = Ctx.CreateClient(options: new CopilotClientOptions { Connection = RuntimeConnection.ForTcp(connectionToken: Token) });

        await GoodClient.StartAsync();
        Port = GoodClient.RuntimePort
            ?? throw new InvalidOperationException("GoodClient is not using TCP mode; RuntimePort is null");
    }

    public async Task DisposeAsync()
    {
        if (GoodClient is not null)
        {
            await GoodClient.ForceStopAsync();
        }

        await Ctx.DisposeAsync();
    }
}

public class ConnectionTokenTests : IClassFixture<ConnectionTokenTestFixture>
{
    private readonly ConnectionTokenTestFixture _fixture;

    public ConnectionTokenTests(ConnectionTokenTestFixture fixture)
    {
        _fixture = fixture;
    }

    [Fact]
    public async Task Connects_With_The_Matching_Token()
    {
        var pong = await _fixture.GoodClient.PingAsync("hi");
        Assert.Equal("pong: hi", pong.Message);
    }

    [Fact]
    public async Task Rejects_A_Wrong_Token()
    {
        var wrongClient = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri($"localhost:{_fixture.Port}", connectionToken: "wrong") });

        try
        {
            var ex = await Assert.ThrowsAnyAsync<Exception>(() => wrongClient.StartAsync());
            Assert.Contains("AUTHENTICATION_FAILED", GetFullMessage(ex));
        }
        finally
        {
            // Best-effort cleanup; ignore stop errors when the client failed to start.
            try { await wrongClient.ForceStopAsync(); } catch (Exception) { }
        }
    }

    [Fact]
    public async Task Rejects_A_Missing_Token_When_One_Is_Required()
    {
        var noTokenClient = new CopilotClient(new CopilotClientOptions { Connection = RuntimeConnection.ForUri($"localhost:{_fixture.Port}") });

        try
        {
            var ex = await Assert.ThrowsAnyAsync<Exception>(() => noTokenClient.StartAsync());
            Assert.Contains("AUTHENTICATION_FAILED", GetFullMessage(ex));
        }
        finally
        {
            // Best-effort cleanup; ignore stop errors when the client failed to start.
            try { await noTokenClient.ForceStopAsync(); } catch (Exception) { }
        }
    }

    private static string GetFullMessage(Exception ex)
    {
        var messages = new List<string>();
        for (var cur = ex; cur is not null; cur = cur.InnerException)
        {
            messages.Add(cur.Message);
        }
        return string.Join(" | ", messages);
    }
}

/// <summary>
/// When the SDK spawns its own CLI in TCP mode without an explicit token, it auto-generates
/// a GUID and round-trips it through the spawned CLI.
/// </summary>
public class ConnectionTokenAutoGeneratedTests : IAsyncLifetime
{
    private E2ETestContext _ctx = null!;
    private CopilotClient _client = null!;

    public async Task InitializeAsync()
    {
        _ctx = await E2ETestContext.CreateAsync();
        _client = _ctx.CreateClient(useStdio: false);
    }

    public async Task DisposeAsync()
    {
        if (_client is not null)
        {
            await _client.ForceStopAsync();
        }

        await _ctx.DisposeAsync();
    }

    [Fact]
    public async Task The_SDK_Auto_Generated_Guid_Round_Trips_Through_The_Spawned_CLI()
    {
        await _client.StartAsync();
        var pong = await _client.PingAsync("hi");
        Assert.Equal("pong: hi", pong.Message);
    }
}