soteria/internal/server/b2_scan_test.go

package server

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sort"
	"strings"
	"testing"
	"time"
)

type fakeS3Object struct {
	key          string
	size         int64
	lastModified time.Time
}

func newFakeS3Server(t *testing.T, buckets []string, objects map[string][]fakeS3Object, missing map[string]bool) *httptest.Server {
	t.Helper()

	sort.Strings(buckets)
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/xml")
		switch {
		case r.URL.Path == "/" && r.URL.RawQuery == "":
			fmt.Fprint(w, `<?xml version="1.0" encoding="UTF-8"?><ListAllMyBucketsResult><Buckets>`)
			for _, bucket := range buckets {
				fmt.Fprintf(w, `<Bucket><Name>%s</Name><CreationDate>2026-04-20T00:00:00Z</CreationDate></Bucket>`, bucket)
			}
			fmt.Fprint(w, `</Buckets></ListAllMyBucketsResult>`)
		case strings.HasSuffix(r.URL.Path, "/") && r.URL.Query().Has("location"):
			fmt.Fprint(w, `<?xml version="1.0" encoding="UTF-8"?><LocationConstraint xmlns="http://s3.amazonaws.com/doc/2006-03-01/">us-west-001</LocationConstraint>`)
		case strings.HasSuffix(r.URL.Path, "/") && r.URL.Query().Get("list-type") == "2":
			bucket := strings.Trim(strings.TrimSuffix(r.URL.Path, "/"), "/")
			if missing[bucket] {
				http.Error(w, "no such bucket", http.StatusNotFound)
				return
			}
			fmt.Fprintf(w, `<?xml version="1.0" encoding="UTF-8"?><ListBucketResult><Name>%s</Name><IsTruncated>false</IsTruncated>`, bucket)
			for _, object := range objects[bucket] {
				fmt.Fprintf(w, `<Contents><Key>%s</Key><LastModified>%s</LastModified><Size>%d</Size></Contents>`,
					object.key,
					object.lastModified.UTC().Format(time.RFC3339),
					object.size,
				)
			}
			fmt.Fprint(w, `</ListBucketResult>`)
		default:
			http.Error(w, "unexpected request", http.StatusNotFound)
		}
	}))
}

func TestScanB2UsageAutoDiscoversBucketsAndAggregatesObjects(t *testing.T) {
	now := time.Now().UTC()
	server := newFakeS3Server(t,
		[]string{"zeta", "alpha"},
		map[string][]fakeS3Object{
			"alpha": {
				{key: "alpha-new", size: 12, lastModified: now.Add(-2 * time.Hour)},
				{key: "alpha-old", size: 30, lastModified: now.Add(-48 * time.Hour)},
			},
			"zeta": {
				{key: "zeta-new", size: 8, lastModified: now.Add(-1 * time.Hour)},
			},
		},
		nil,
	)
	defer server.Close()

	result, err := scanB2Usage(context.Background(), b2Credentials{
		Endpoint:        server.URL,
		Region:          "us-west-001",
		AccessKeyID:     "atlas-key",
		SecretAccessKey: "atlas-secret",
	}, nil)
	if err != nil {
		t.Fatalf("scan b2 usage autodiscovery: %v", err)
	}

	if !result.Enabled || !result.Available {
		t.Fatalf("expected available B2 usage result, got %#v", result)
	}
	if len(result.Buckets) != 2 || result.Buckets[0].Name != "alpha" || result.Buckets[1].Name != "zeta" {
		t.Fatalf("expected sorted bucket results, got %#v", result.Buckets)
	}
	if result.TotalObjects != 3 || result.TotalBytes != 50 {
		t.Fatalf("expected aggregated object totals, got %#v", result)
	}
	if result.RecentObjects24h != 2 || result.RecentBytes24h != 20 {
		t.Fatalf("expected recent object totals, got %#v", result)
	}
	if result.Buckets[0].RecentObjects24h != 1 || result.Buckets[0].RecentBytes24h != 12 || result.Buckets[0].LastModifiedAt == "" {
		t.Fatalf("expected alpha bucket recent stats, got %#v", result.Buckets[0])
	}
}

func TestScanB2UsageConfiguredBucketsAndErrorBranches(t *testing.T) {
	now := time.Now().UTC()
	server := newFakeS3Server(t,
		[]string{},
		map[string][]fakeS3Object{
			"alpha": {{key: "alpha", size: 1, lastModified: now.Add(-1 * time.Hour)}},
			"beta":  {{key: "beta", size: 2, lastModified: now.Add(-2 * time.Hour)}},
		},
		map[string]bool{"missing": true},
	)
	defer server.Close()

	result, err := scanB2Usage(context.Background(), b2Credentials{
		Endpoint:        server.URL,
		Region:          "us-west-001",
		AccessKeyID:     "atlas-key",
		SecretAccessKey: "atlas-secret",
	}, []string{" beta ", "", "alpha"})
	if err != nil {
		t.Fatalf("scan b2 usage configured buckets: %v", err)
	}
	if len(result.Buckets) != 2 || result.Buckets[0].Name != "alpha" || result.Buckets[1].Name != "beta" {
		t.Fatalf("expected configured buckets to be trimmed and sorted, got %#v", result.Buckets)
	}

	emptyServer := newFakeS3Server(t, nil, nil, nil)
	defer emptyServer.Close()
	if _, err := scanB2Usage(context.Background(), b2Credentials{
		Endpoint:        emptyServer.URL,
		Region:          "us-west-001",
		AccessKeyID:     "atlas-key",
		SecretAccessKey: "atlas-secret",
	}, nil); err == nil || !strings.Contains(err.Error(), "no B2 buckets available for scan") {
		t.Fatalf("expected no-buckets error, got %v", err)
	}

	if _, err := scanB2Usage(context.Background(), b2Credentials{
		Endpoint:        server.URL,
		Region:          "us-west-001",
		AccessKeyID:     "atlas-key",
		SecretAccessKey: "atlas-secret",
	}, []string{"missing"}); err == nil || !strings.Contains(err.Error(), "scan B2 bucket missing") {
		t.Fatalf("expected bucket scan error, got %v", err)
	}
}