prefab-page-detail/Prefab.Tests/Infrastructure/AspireInfraHost.cs

using Microsoft.Extensions.Configuration;
using Microsoft.Extensions.Hosting;

namespace Prefab.Tests.Infrastructure;

/// <summary>
/// Starts the Aspire AppHost while keeping only infrastructure resources (e.g., SQL, cache).
/// </summary>
public sealed class AspireInfraHost : IAsyncLifetime
{
    private readonly PrefabHarnessOptions _options;
    private readonly string[] _infraResourcesToKeep;
    private DistributedApplication? _app;
    private readonly Dictionary<string, ProjectResource> _projects = new(StringComparer.OrdinalIgnoreCase);

    public AspireInfraHost(PrefabHarnessOptions options, params string[] infraResourcesToKeep)
    {
        _options = options;
        _infraResourcesToKeep = infraResourcesToKeep;
    }

    public DistributedApplication App => _app ?? throw new InvalidOperationException("Aspire is not running.");
    public IServiceProvider Services => App.Services;

    public async ValueTask InitializeAsync()
    {
        var appArgs = new List<string>();
        if (_options.Mode == RunMode.EphemeralIsolated)
        {
            var containerName = $"prefab-sql-{Guid.NewGuid():N}"[..19];
            appArgs.Add("SqlServer:UseDataVolume=false");
            appArgs.Add("SqlServer:HostPort=");
            appArgs.Add($"SqlServer:ContainerName={containerName}");
        }
        if (_options.DisablePortRandomization)
        {
            appArgs.Add("DcpPublisher:RandomizePorts=false");
        }

        var builder = await DistributedApplicationTestingBuilder.CreateAsync<Projects.Prefab_AppHost>(
            appArgs.ToArray(),
            configureBuilder: (appOptions, hostSettings) =>
            {
                appOptions.DisableDashboard = !_options.EnableAspireDashboard;

                hostSettings ??= new HostApplicationBuilderSettings();
                hostSettings.Configuration ??= new ConfigurationManager();
            });

        if (_options.DisablePortRandomization)
        {
            builder.Configuration["DcpPublisher:RandomizePorts"] = "false";
        }
        builder.Configuration["ASPIRE_ALLOW_UNSECURED_TRANSPORT"] = "true";

        foreach (var project in builder.Resources.OfType<ProjectResource>())
        {
            _projects[project.Name] = project;
        }

        var keep = new HashSet<string>(_infraResourcesToKeep, StringComparer.OrdinalIgnoreCase);

        // Ensure parent resources remain so infra dependencies stay intact.
        var added = 0;
        do
        {
            added = 0;
            foreach (var resource in builder.Resources)
            {
                if (keep.Contains(resource.Name))
                {
                    if (resource is IResourceWithParent withParent && withParent.Parent is { } parent && keep.Add(parent.Name))
                    {
                        added++;
                    }

                    var relationships = resource.Annotations.OfType<ResourceRelationshipAnnotation>();
                    foreach (var relationship in relationships.Where(r => r.Type == "Parent"))
                    {
                        if (keep.Add(relationship.Resource.Name))
                        {
                            added++;
                        }
                    }
                }
            }
        } while (added > 0);

        foreach (var resource in builder.Resources.Where(r => !keep.Contains(r.Name)).ToArray())
        {
            builder.Resources.Remove(resource);
        }

        if (_options.Mode == RunMode.EphemeralIsolated)
        {
            foreach (var resource in builder.Resources)
            {
                var lifetime = resource.Annotations.OfType<ContainerLifetimeAnnotation>().FirstOrDefault();
                if (lifetime is not null)
                {
                    resource.Annotations.Remove(lifetime);
                }
            }
        }

        _app = await builder.BuildAsync();
        await _app.StartAsync();

        foreach (var resource in builder.Resources)
        {
            await _app.ResourceNotifications.WaitForResourceHealthyAsync(resource.Name);
        }
    }

    public async ValueTask DisposeAsync()
    {
        if (_app is not null)
        {
            await _app.DisposeAsync();
            _app = null;
        }
    }

    public ProjectResource GetProject(string name)
    {
        if (_projects.TryGetValue(name, out var project))
        {
            return project;
        }

        throw new ArgumentException($"Project '{name}' not found in AppHost model.", nameof(name));
    }

    public async Task<string> GetConnectionStringAsync(string resourceName, CancellationToken cancellationToken = default) =>
        (await App.GetConnectionStringAsync(resourceName, cancellationToken))
        ?? throw new InvalidOperationException($"Connection string for '{resourceName}' is not available.");

    public static async Task<Dictionary<string, string?>> ResolveConfigForProjectAsync(
        ProjectResource project,
        IServiceProvider services,
        CancellationToken cancellationToken = default)
    {
        var config = new Dictionary<string, string?>(StringComparer.OrdinalIgnoreCase);

        if (project is not IResourceWithEnvironment withEnvironment ||
            !withEnvironment.TryGetEnvironmentVariables(out var annotations))
        {
            return config;
        }

        var execOptions = new DistributedApplicationExecutionContextOptions(DistributedApplicationOperation.Run)
        {
            ServiceProvider = services
        };
        var execContext = new DistributedApplicationExecutionContext(execOptions);
        var envContext = new EnvironmentCallbackContext(execContext, cancellationToken: cancellationToken);

        foreach (var annotation in annotations)
        {
            await annotation.Callback(envContext);
        }

        foreach (var (key, value) in envContext.EnvironmentVariables)
        {
            if (string.Equals(key, "ASPNETCORE_URLS", StringComparison.OrdinalIgnoreCase))
            {
                continue;
            }

            var resolved = value switch
            {
                string s => s,
                IValueProvider provider => await SafeGet(provider, cancellationToken),
                _ => null
            };

            if (resolved is not null)
            {
                config[key.Replace("__", ":")] = resolved;
            }
        }

        return config;

        static async Task<string?> SafeGet(IValueProvider provider, CancellationToken ct)
        {
            try
            {
                return await provider.GetValueAsync(ct);
            }
            catch
            {
                return null;
            }
        }
    }
}