git.stella-ops.org/src/Scheduler/__Tests/StellaOps.Scheduler.ImpactIndex.Tests/RoaringImpactIndexTests.cs

using System;
using System.Collections.Immutable;
using System.IO;
using FluentAssertions;
using Microsoft.Extensions.Logging.Abstractions;
using StellaOps.Scheduler.ImpactIndex.Ingestion;
using StellaOps.Scheduler.Models;
using StellaOps.Scanner.Core.Contracts;
using StellaOps.Scanner.Emit.Index;
using Xunit;

namespace StellaOps.Scheduler.ImpactIndex.Tests;

public sealed class RoaringImpactIndexTests
{
    [Fact]
    public async Task IngestAsync_RegistersComponentsAndUsage()
    {
        var (stream, digest) = CreateBomIndex(
            ComponentIdentity.Create("pkg:npm/a@1.0.0", "a", "1.0.0", "pkg:npm/a@1.0.0"),
            ComponentUsage.Create(true, new[] { "/app/start.sh" }));

        var index = new RoaringImpactIndex(NullLogger<RoaringImpactIndex>.Instance);
        var request = new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-alpha",
            ImageDigest = digest,
            Registry = "docker.io",
            Repository = "library/alpine",
            Namespaces = ImmutableArray.Create("team-a"),
            Tags = ImmutableArray.Create("3.20"),
            Labels = ImmutableSortedDictionary.CreateRange(StringComparer.OrdinalIgnoreCase, new[]
            {
                new KeyValuePair<string, string>("env", "prod")
            }),
            BomIndexStream = stream,
        };

        await index.IngestAsync(request, CancellationToken.None);

        var selector = new Selector(SelectorScope.AllImages, tenantId: "tenant-alpha");
        var impactSet = await index.ResolveByPurlsAsync(new[] { "pkg:npm/a@1.0.0" }, usageOnly: false, selector);

        impactSet.Images.Should().HaveCount(1);
        impactSet.Images[0].ImageDigest.Should().Be(digest);
        impactSet.Images[0].Tags.Should().ContainSingle(tag => tag == "3.20");
        impactSet.Images[0].UsedByEntrypoint.Should().BeTrue();

        var usageOnly = await index.ResolveByPurlsAsync(new[] { "pkg:npm/a@1.0.0" }, usageOnly: true, selector);
        usageOnly.Images.Should().HaveCount(1);
    }

    [Fact]
    public async Task IngestAsync_ReplacesExistingImageData()
    {
        var component = ComponentIdentity.Create("pkg:npm/a@1.0.0", "a", "1.0.0", "pkg:npm/a@1.0.0");
        var (initialStream, digest) = CreateBomIndex(component, ComponentUsage.Create(false));
        var index = new RoaringImpactIndex(NullLogger<RoaringImpactIndex>.Instance);

        await index.IngestAsync(new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-alpha",
            ImageDigest = digest,
            Registry = "docker.io",
            Repository = "library/alpine",
            Tags = ImmutableArray.Create("v1"),
            BomIndexStream = initialStream,
        });

        var (updatedStream, _) = CreateBomIndex(component, ComponentUsage.Create(true, new[] { "/start.sh" }), digest);
        await index.IngestAsync(new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-alpha",
            ImageDigest = digest,
            Registry = "docker.io",
            Repository = "library/alpine",
            Tags = ImmutableArray.Create("v2"),
            BomIndexStream = updatedStream,
        });

        var selector = new Selector(SelectorScope.AllImages, tenantId: "tenant-alpha");
        var impactSet = await index.ResolveByPurlsAsync(new[] { "pkg:npm/a@1.0.0" }, usageOnly: true, selector);

        impactSet.Images.Should().HaveCount(1);
        impactSet.Images[0].Tags.Should().ContainSingle(tag => tag == "v2");
        impactSet.Images[0].UsedByEntrypoint.Should().BeTrue();
    }

    [Fact]
    public async Task ResolveByPurlsAsync_RespectsTenantNamespaceAndTagFilters()
    {
        var component = ComponentIdentity.Create("pkg:npm/a@1.0.0", "a", "1.0.0", "pkg:npm/a@1.0.0");
        var (tenantStream, tenantDigest) = CreateBomIndex(component, ComponentUsage.Create(true, new[] { "/start.sh" }));
        var (otherStream, otherDigest) = CreateBomIndex(component, ComponentUsage.Create(false));

        var index = new RoaringImpactIndex(NullLogger<RoaringImpactIndex>.Instance);

        await index.IngestAsync(new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-alpha",
            ImageDigest = tenantDigest,
            Registry = "docker.io",
            Repository = "library/service",
            Namespaces = ImmutableArray.Create("team-alpha"),
            Tags = ImmutableArray.Create("prod-eu"),
            BomIndexStream = tenantStream,
        });

        await index.IngestAsync(new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-beta",
            ImageDigest = otherDigest,
            Registry = "docker.io",
            Repository = "library/service",
            Namespaces = ImmutableArray.Create("team-beta"),
            Tags = ImmutableArray.Create("staging-us"),
            BomIndexStream = otherStream,
        });

        var selector = new Selector(
            SelectorScope.AllImages,
            tenantId: "tenant-alpha",
            namespaces: new[] { "team-alpha" },
            includeTags: new[] { "prod-*" });

        var result = await index.ResolveByPurlsAsync(new[] { "pkg:npm/a@1.0.0" }, usageOnly: true, selector);

        result.Images.Should().ContainSingle(image => image.ImageDigest == tenantDigest);
        result.Images[0].Tags.Should().Contain("prod-eu");
    }

    [Fact]
    public async Task ResolveAllAsync_UsageOnlyFiltersEntrypointImages()
    {
        var component = ComponentIdentity.Create("pkg:npm/a@1.0.0", "a", "1.0.0", "pkg:npm/a@1.0.0");
        var (entryStream, entryDigest) = CreateBomIndex(component, ComponentUsage.Create(true, new[] { "/start.sh" }));
        var nonEntryDigestValue = "sha256:" + new string('1', 64);
        var (nonEntryStream, nonEntryDigest) = CreateBomIndex(component, ComponentUsage.Create(false), nonEntryDigestValue);

        var index = new RoaringImpactIndex(NullLogger<RoaringImpactIndex>.Instance);

        await index.IngestAsync(new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-alpha",
            ImageDigest = entryDigest,
            Registry = "docker.io",
            Repository = "library/service",
            BomIndexStream = entryStream,
        });

        await index.IngestAsync(new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-alpha",
            ImageDigest = nonEntryDigest,
            Registry = "docker.io",
            Repository = "library/service",
            BomIndexStream = nonEntryStream,
        });

        var selector = new Selector(SelectorScope.AllImages, tenantId: "tenant-alpha");

        var usageOnly = await index.ResolveAllAsync(selector, usageOnly: true);
        usageOnly.Images.Should().ContainSingle(image => image.ImageDigest == entryDigest);

        var allImages = await index.ResolveAllAsync(selector, usageOnly: false);
        allImages.Images.Should().HaveCount(2);
    }

    [Fact]
    public async Task RemoveAsync_RemovesImageAndComponents()
    {
        var component = ComponentIdentity.Create("pkg:npm/a@1.0.0", "a", "1.0.0", "pkg:npm/a@1.0.0");
        var (stream1, digest1) = CreateBomIndex(component, ComponentUsage.Create(true, new[] { "/start.sh" }));
        var (stream2, digest2) = CreateBomIndex(component, ComponentUsage.Create(true, new[] { "/start.sh" }));

        var index = new RoaringImpactIndex(NullLogger<RoaringImpactIndex>.Instance);

        await index.IngestAsync(new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-alpha",
            ImageDigest = digest1,
            Registry = "docker.io",
            Repository = "library/service",
            BomIndexStream = stream1,
        });

        await index.IngestAsync(new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-alpha",
            ImageDigest = digest2,
            Registry = "docker.io",
            Repository = "library/service",
            BomIndexStream = stream2,
        });

        await index.RemoveAsync(digest1);

        var selector = new Selector(SelectorScope.AllImages, tenantId: "tenant-alpha");
        var impact = await index.ResolveByPurlsAsync(new[] { "pkg:npm/a@1.0.0" }, usageOnly: false, selector);

        impact.Images.Should().ContainSingle(img => img.ImageDigest == digest2);
    }

    [Fact]
    public async Task CreateSnapshotAsync_CompactsIdsAndRestores()
    {
        var component = ComponentIdentity.Create("pkg:npm/a@1.0.0", "a", "1.0.0", "pkg:npm/a@1.0.0");
        var (stream1, digest1) = CreateBomIndex(component, ComponentUsage.Create(true, new[] { "/start.sh" }));
        var (stream2, digest2) = CreateBomIndex(component, ComponentUsage.Create(false));

        var index = new RoaringImpactIndex(NullLogger<RoaringImpactIndex>.Instance);

        await index.IngestAsync(new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-alpha",
            ImageDigest = digest1,
            Registry = "docker.io",
            Repository = "library/service",
            BomIndexStream = stream1,
        });

        await index.IngestAsync(new ImpactIndexIngestionRequest
        {
            TenantId = "tenant-alpha",
            ImageDigest = digest2,
            Registry = "docker.io",
            Repository = "library/service",
            BomIndexStream = stream2,
        });

        await index.RemoveAsync(digest1);

        var snapshot = await index.CreateSnapshotAsync();
        snapshot.SnapshotId.Should().MatchRegex("^snap-[0-9a-f]{64}$");

        var restored = new RoaringImpactIndex(NullLogger<RoaringImpactIndex>.Instance);
        await restored.RestoreSnapshotAsync(snapshot);

        var selector = new Selector(SelectorScope.AllImages, tenantId: "tenant-alpha");
        var resolved = await restored.ResolveAllAsync(selector, usageOnly: false);

        resolved.Images.Should().ContainSingle(img => img.ImageDigest == digest2);
        resolved.SnapshotId.Should().Be(snapshot.SnapshotId);
    }

    private static (Stream Stream, string Digest) CreateBomIndex(ComponentIdentity identity, ComponentUsage usage, string? digest = null)
    {
        var layer = LayerComponentFragment.Create(
            "sha256:layer1",
            new[]
            {
                new ComponentRecord
                {
                    Identity = identity,
                    LayerDigest = "sha256:layer1",
                    Usage = usage,
                }
            });

        var graph = ComponentGraphBuilder.Build(new[] { layer });
        var effectiveDigest = digest ?? "sha256:" + Guid.NewGuid().ToString("N");
        var builder = new BomIndexBuilder();
        var artifact = builder.Build(new BomIndexBuildRequest
        {
            ImageDigest = effectiveDigest,
            Graph = graph,
            GeneratedAt = DateTimeOffset.UtcNow,
        });

        return (new MemoryStream(artifact.Bytes, writable: false), effectiveDigest);
    }
}