package enrichment

import (
	"context"
	"testing"

	"github.com/tidbcloud/infra-provider/regional-server/pkg/metrics/objectstorage/provider"
)

type staticSource map[string]BucketLabels

func (s staticSource) Load(context.Context) (map[string]BucketLabels, error) { return s, nil }

func TestCollectorEnrichesMatchedBucketsAndLeavesMissingBucketsEmpty(t *testing.T) {
	base := collectorFunc(func(context.Context) (provider.Snapshot, error) {
		return provider.Snapshot{Samples: []provider.BucketSample{
			{Bucket: "managed", StorageBytes: 1},
			{Bucket: "unmanaged", StorageBytes: 2},
		}}, nil
	})
	collector := NewCollector(base, staticSource{
		"managed": {TenantID: "tenant-1", SharedPoolID: "pool-1"},
	})

	snapshot, err := collector.Refresh(context.Background())
	if err != nil {
		t.Fatalf("Refresh() error = %v", err)
	}
	if snapshot.Samples[0].TenantID != "tenant-1" || snapshot.Samples[0].SharedPoolID != "pool-1" {
		t.Fatalf("managed sample not enriched: %+v", snapshot.Samples[0])
	}
	if snapshot.Samples[1].TenantID != "" || snapshot.Samples[1].SharedPoolID != "" {
		t.Fatalf("unmanaged sample should keep empty labels: %+v", snapshot.Samples[1])
	}
}

type collectorFunc func(context.Context) (provider.Snapshot, error)

func (f collectorFunc) Name() string                                           { return "base" }
func (f collectorFunc) Refresh(ctx context.Context) (provider.Snapshot, error) { return f(ctx) }