mesh-catalog/modules/claude-licence-manager/cmd/claude-licence-manager/store_test.go

package main

import (
	"context"
	"os"
	"strings"
	"testing"
	"time"
)

// Against a real postgres, because the questions are the database's: does the schema apply twice, does a
// lease refuse a second holder, does a generation only grow. Skipped unless one is named:
//
//	docker run -d --rm --name licmgr-pg -e POSTGRES_PASSWORD=t -p 15498:5432 postgres:16-alpine
//	MESH_TEST_POSTGRES=postgres://postgres:t@127.0.0.1:15498/postgres go test ./...
func TestTheStoreOnPostgres(t *testing.T) {
	url := os.Getenv("MESH_TEST_POSTGRES")
	if url == "" {
		t.Skip("MESH_TEST_POSTGRES unset")
	}
	ctx := context.Background()
	s, err := OpenPgStore(ctx, url)
	if err != nil {
		t.Fatal(err)
	}
	defer s.Close()
	for _, table := range []string{"binding", "licence", "lease", "offered", "usage", "audit"} {
		_, _ = s.pool.Exec(ctx, "drop table if exists "+table+" cascade")
	}
	_, _ = s.pool.Exec(ctx, "drop sequence if exists binding_generation")
	for i := 0; i < 2; i++ {
		if err := s.Migrate(ctx); err != nil {
			t.Fatalf("migration %d: %v", i+1, err)
		}
	}
	l := Licence{Name: "a@example.org", Kind: "subscription", AccountUUID: "u-1", Email: "a@example.org", Sealed: "v1.x.y",
		RefreshFingerprint: "sha256:1", AccessExpiresAt: 1, RefreshExpiresAt: 2, AdoptedAt: 3}
	if err := s.SaveLicence(ctx, l); err != nil {
		t.Fatal(err)
	}
	l.Failures = 2
	_ = s.SaveLicence(ctx, l)
	if got, _ := s.LicenceForAccount(ctx, "u-1"); got == nil || got.Failures != 2 || got.OrganizationUUID != "" {
		t.Fatalf("%+v", got)
	}
	if none, err := s.Licence(ctx, "nobody"); none != nil || err != nil {
		t.Fatalf("%v %v", none, err)
	}
	lease := func(holder string) bool {
		ok, err := s.Lease(ctx, "licence:a", holder, time.Minute)
		if err != nil {
			t.Fatal(err)
		}
		return ok
	}
	if !lease("one") || lease("two") || !lease("one") {
		t.Fatal("a lease did not refuse a second holder, or its own holder could not renew it")
	}
	_ = s.Unlease(ctx, "licence:a", "one")
	if !lease("two") {
		t.Fatal("a released lease was not taken")
	}
	b1, _ := s.Bind(ctx, "laptop", l.Name)
	adv, _ := s.Advance(ctx, l.Name)
	b3, _ := s.Bind(ctx, "laptop", l.Name)
	if len(adv) != 1 || !(b1.Generation < adv[0].Generation && adv[0].Generation < b3.Generation) {
		t.Fatalf("generations %d %v %d", b1.Generation, adv, b3.Generation)
	}
	_ = s.RecordOutcome(ctx, "sha256:x", "laptop", "u-1", Dead, "400")
	if o, _ := s.Outcome(ctx, "sha256:x"); o != Dead {
		t.Fatalf("outcome %q", o)
	}
	if o, _ := s.Outcome(ctx, "sha256:none"); o != "" {
		t.Fatalf("an unknown login is %q", o)
	}
	pct := 1.0
	_ = s.RecordUsage(ctx, l.Name, 5, UsageReading{SessionPct: &pct}, map[string]any{})
	if u, _ := s.Usage(ctx, "", 5); len(u) != 1 || *u[0].Reading.SessionPct != 1 {
		t.Fatalf("usage %v", u)
	}
	if err := s.Audit(ctx, "bound", map[string]any{"consumer": "laptop"}); err != nil {
		t.Fatal(err)
	}
	if ok, _ := s.Unbind(ctx, "laptop"); !ok {
		t.Fatal("unbind")
	}
	all, _ := s.Licences(ctx)
	if len(all) != 1 || !strings.HasPrefix(all[0].Sealed, "v1.") {
		t.Fatalf("%v", all)
	}
}

func TestARefreshThatDoesNotRotateKeepsTheRefreshToken(t *testing.T) {
	prev := FullGrant{AccessToken: "a", RefreshToken: "r", ExpiresAt: 1}
	in := int64(60)
	kept := NextGrant(prev, tokenResponse{AccessToken: "a2", ExpiresIn: &in}, 1000)
	if !kept.OK || kept.Grant.RefreshToken != "r" || kept.Grant.ExpiresAt != 61_000 {
		t.Fatalf("%+v", kept)
	}
	rotated := NextGrant(prev, tokenResponse{AccessToken: "a3", RefreshToken: "r2"}, 0)
	if rotated.Grant.RefreshToken != "r2" {
		t.Fatalf("%+v", rotated)
	}
	if NextGrant(prev, tokenResponse{}, 0).OK {
		t.Fatal("an answer with no access token was a grant")
	}
}

func TestAGrantAtRestOpensOnlyWithItsKey(t *testing.T) {
	a, _ := NewCrypt("one key")
	b, _ := NewCrypt("another key")
	sealed := a.Seal(`{"refreshToken":"rt"}`)
	if strings.Contains(sealed, "rt") {
		t.Fatal("stored in the clear")
	}
	if got, err := a.Open(sealed); err != nil || got != `{"refreshToken":"rt"}` {
		t.Fatalf("%q %v", got, err)
	}
	if _, err := b.Open(sealed); err == nil {
		t.Fatal("opened with another key")
	}
	if _, err := NewCrypt("  "); err == nil {
		t.Fatal("an empty key was accepted")
	}
}