MeterVault/tests/Integration.Tests/Ingestion/HomeAssistantWebSocketWorkerTests.cs

using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using MeterVault.Core.Domain;
using MeterVault.Infrastructure.Ingestion;
using MeterVault.Infrastructure.Persistence;
using Microsoft.AspNetCore.Builder;
using Microsoft.AspNetCore.Hosting;
using Microsoft.AspNetCore.Hosting.Server;
using Microsoft.AspNetCore.Hosting.Server.Features;
using Microsoft.AspNetCore.Http;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Logging.Abstractions;

namespace MeterVault.Integration.Tests.Ingestion;

/// <summary>
/// End-to-end proof of the HA WebSocket push path against an in-process fake Home Assistant server:
/// the worker performs the auth handshake, subscribes to <c>state_changed</c>, and a pushed change
/// for a configured entity lands as a <c>reading</c>. This exercises the real ClientWebSocket
/// transport and handshake sequencing that the pure-protocol unit tests can't.
/// </summary>
[Collection("Timescale")]
public sealed class HomeAssistantWebSocketWorkerTests(TimescaleFixture fx)
{
    private const string TokenEnvVar = "MV_TEST_HA_WS_TOKEN";

    [Fact]
    public async Task Pushed_state_change_becomes_a_reading()
    {
        await using var db = fx.CreateContext();
        Environment.SetEnvironmentVariable(TokenEnvVar, "test-token");
        await using var fake = await FakeHaServer.StartAsync(entityId: "sensor.house_power", state: "4711");

        try
        {
            var type = new EnergyType { Key = "ha_ws_test", DisplayName = "HA WS", BaseUnit = "kWh", DefaultMode = MeterMode.CumulativeCounter };
            db.EnergyTypes.Add(type);
            await db.SaveChangesAsync();

            var meter = new Meter { Name = "HA WS Meter", EnergyTypeId = type.Id, Mode = MeterMode.CumulativeCounter, Unit = "kWh" };
            db.Meters.Add(meter);
            await db.SaveChangesAsync();

            var endpoint = new IngestionEndpoint
            {
                Type = EndpointType.HomeAssistant,
                Name = "Fake HA",
                IsEnabled = true,
                Config = new HaEndpointConfig { BaseUrl = fake.BaseUrl, TokenEnv = TokenEnvVar, UseWebSocket = true }.ToJson(),
            };
            db.IngestionEndpoints.Add(endpoint);
            await db.SaveChangesAsync();

            db.MeterSources.Add(new MeterSource
            {
                MeterId = meter.Id,
                SourceType = SourceType.HomeAssistant,
                EndpointId = endpoint.Id,
                IsEnabled = true,
                Config = JsonSerializer.Serialize(new { entityId = "sensor.house_power" }),
            });
            await db.SaveChangesAsync();

            await using var provider = BuildProvider(fx.ConnectionString);
            var worker = new HomeAssistantWebSocketWorker(
                provider.GetRequiredService<IServiceScopeFactory>(), NullLogger<HomeAssistantWebSocketWorker>.Instance);

            await worker.StartAsync(CancellationToken.None);
            try
            {
                Reading? reading = null;
                for (var i = 0; i < 60 && reading is null; i++)
                {
                    await Task.Delay(200);
                    reading = await db.Readings.AsNoTracking().FirstOrDefaultAsync(r => r.MeterId == meter.Id);
                }

                Assert.NotNull(reading);
                Assert.Equal(4711, reading!.Value, 3);
                Assert.Equal(ReadingQuality.Measured, reading.Quality);
            }
            finally
            {
                await worker.StopAsync(CancellationToken.None);
            }
        }
        finally
        {
            Environment.SetEnvironmentVariable(TokenEnvVar, null);
            await db.Readings.ExecuteDeleteAsync();
            await db.MeterSources.ExecuteDeleteAsync();
            await db.IngestionEndpoints.ExecuteDeleteAsync();
            await db.Meters.ExecuteDeleteAsync();
            await db.EnergyTypes.Where(t => t.Key == "ha_ws_test").ExecuteDeleteAsync();
        }
    }

    private static ServiceProvider BuildProvider(string connectionString)
    {
        var services = new ServiceCollection();
        services.AddLogging();
        services.AddDbContextFactory<MeterVaultDbContext>(o => o
            .UseNpgsql(connectionString, n => n.MigrationsAssembly(typeof(MeterVaultDbContext).Assembly.FullName))
            .UseSnakeCaseNamingConvention());
        services.AddScoped<MeterVaultDbContext>(sp => sp.GetRequiredService<IDbContextFactory<MeterVaultDbContext>>().CreateDbContext());
        services.AddScoped<IngestionService>();
        return services.BuildServiceProvider();
    }

    /// <summary>A minimal Home Assistant WebSocket server: handshake, then push one state_changed event.</summary>
    private sealed class FakeHaServer : IAsyncDisposable
    {
        private readonly WebApplication _app;

        private FakeHaServer(WebApplication app, string baseUrl)
        {
            _app = app;
            BaseUrl = baseUrl;
        }

        public string BaseUrl { get; }

        public static async Task<FakeHaServer> StartAsync(string entityId, string state)
        {
            var builder = WebApplication.CreateBuilder();
            builder.Logging.ClearProviders();
            builder.WebHost.UseUrls("http://127.0.0.1:0");
            var app = builder.Build();
            app.UseWebSockets();

            var eventJson =
                "{\"id\":1,\"type\":\"event\",\"event\":{\"event_type\":\"state_changed\",\"data\":{\"entity_id\":\""
                + entityId + "\",\"new_state\":{\"entity_id\":\"" + entityId + "\",\"state\":\"" + state
                + "\",\"attributes\":{},\"last_updated\":\"2024-06-15T10:00:00+00:00\"}}}}";

            app.Map("/api/websocket", async context =>
            {
                if (!context.WebSockets.IsWebSocketRequest)
                {
                    context.Response.StatusCode = 400;
                    return;
                }

                using var ws = await context.WebSockets.AcceptWebSocketAsync();
                await SendAsync(ws, """{"type":"auth_required","ha_version":"2024.6"}""");
                await ReceiveAsync(ws); // client "auth"
                await SendAsync(ws, """{"type":"auth_ok"}""");
                await ReceiveAsync(ws); // client "subscribe_events"
                await SendAsync(ws, """{"id":1,"type":"result","success":true}""");
                await SendAsync(ws, eventJson);

                try
                {
                    await Task.Delay(Timeout.Infinite, context.RequestAborted);
                }
                catch (OperationCanceledException)
                {
                    // client/test closed — expected.
                }
            });

            await app.StartAsync();
            var address = app.Services.GetRequiredService<IServer>().Features
                .Get<IServerAddressesFeature>()!.Addresses.First();
            return new FakeHaServer(app, address);
        }

        public async ValueTask DisposeAsync() => await _app.DisposeAsync();

        private static Task SendAsync(WebSocket ws, string json) =>
            ws.SendAsync(Encoding.UTF8.GetBytes(json), WebSocketMessageType.Text, endOfMessage: true, CancellationToken.None);

        private static async Task ReceiveAsync(WebSocket ws)
        {
            var buffer = new byte[8192];
            await ws.ReceiveAsync(buffer, CancellationToken.None);
        }
    }
}