pangolin/server/internal/provision/fakes_test.go

package provision

import (
	"context"
	"fmt"
	"sync"
	"time"
)

// ---------------------------------------------------------------------------
// fakeStore — an in-memory Store that records the ordered sequence of status
// transitions so tests can assert make-before-break timing.
// ---------------------------------------------------------------------------

type statusTransition struct {
	Seq    int
	NodeID int64
	Status Status
}

type fakeStore struct {
	mu sync.Mutex

	nodes     map[int64]*Node
	byUUID    map[string]int64
	providers map[int64]*Provider
	idem      map[string]string
	replaces  map[string]*Replacement
	events    []struct {
		NodeID int64
		Event  Event
		Detail string
	}
	audits      []string
	transitions []statusTransition

	nextNodeID int64
	version    int64
	seq        int
}

func newFakeStore() *fakeStore {
	return &fakeStore{
		nodes:      map[int64]*Node{},
		byUUID:     map[string]int64{},
		providers:  map[int64]*Provider{},
		idem:       map[string]string{},
		replaces:   map[string]*Replacement{},
		nextNodeID: 0,
		version:    1,
	}
}

func (f *fakeStore) addProvider(p *Provider) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.providers[p.ID] = p
}

func (f *fakeStore) InsertNode(_ context.Context, n *Node) (int64, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.nextNodeID++
	id := f.nextNodeID
	cp := *n
	cp.ID = id
	if cp.Status == "" {
		cp.Status = StatusProvisioning
	}
	cp.CreatedAt = time.Now()
	f.nodes[id] = &cp
	f.byUUID[cp.UUID] = id
	f.seq++
	f.transitions = append(f.transitions, statusTransition{f.seq, id, cp.Status})
	return id, nil
}

func (f *fakeStore) GetNode(_ context.Context, id int64) (*Node, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	n, ok := f.nodes[id]
	if !ok {
		return nil, nil
	}
	cp := *n
	return &cp, nil
}

func (f *fakeStore) GetNodeByUUID(_ context.Context, uuid string) (*Node, error) {
	f.mu.Lock()
	id, ok := f.byUUID[uuid]
	f.mu.Unlock()
	if !ok {
		return nil, nil
	}
	return f.GetNode(context.Background(), id)
}

func (f *fakeStore) UpdateNodeStatus(_ context.Context, id int64, status Status) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	n, ok := f.nodes[id]
	if !ok {
		return fmt.Errorf("fakeStore: node %d not found", id)
	}
	n.Status = status
	f.seq++
	f.transitions = append(f.transitions, statusTransition{f.seq, id, status})
	return nil
}

func (f *fakeStore) UpdateNodeEndpoint(_ context.Context, id int64, endpoint string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if n, ok := f.nodes[id]; ok {
		n.Endpoint = endpoint
	}
	return nil
}

func (f *fakeStore) SetNodeInstance(_ context.Context, id int64, instanceID, endpoint string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if n, ok := f.nodes[id]; ok {
		n.ProviderInstanceID = instanceID
		n.Endpoint = endpoint
	}
	return nil
}

func (f *fakeStore) SetNodeWeight(_ context.Context, id int64, weight int) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if n, ok := f.nodes[id]; ok {
		n.Weight = weight
	}
	return nil
}

func (f *fakeStore) ListNodesByPool(_ context.Context, pool Pool, status Status) ([]*Node, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	var out []*Node
	for _, n := range f.nodes {
		p := f.providers[n.ProviderID]
		if p == nil || p.Pool != pool {
			continue
		}
		if status != "" && n.Status != status {
			continue
		}
		cp := *n
		out = append(out, &cp)
	}
	return out, nil
}

func (f *fakeStore) ListProviders(_ context.Context, pool Pool) ([]*Provider, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	var out []*Provider
	for _, p := range f.providers {
		if !p.Enabled {
			continue
		}
		if pool != "" && p.Pool != pool {
			continue
		}
		cp := *p
		out = append(out, &cp)
	}
	return out, nil
}

func (f *fakeStore) GetProvider(_ context.Context, id int64) (*Provider, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	p, ok := f.providers[id]
	if !ok {
		return nil, nil
	}
	cp := *p
	return &cp, nil
}

func (f *fakeStore) WriteNodeEvent(_ context.Context, nodeID int64, event Event, detail string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.events = append(f.events, struct {
		NodeID int64
		Event  Event
		Detail string
	}{nodeID, event, detail})
	return nil
}

func (f *fakeStore) WriteAuditLog(_ context.Context, actor, action, target, meta string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.audits = append(f.audits, fmt.Sprintf("%s|%s|%s", actor, action, target))
	return nil
}

func (f *fakeStore) BumpDirectoryVersion(_ context.Context) (int64, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.version++
	return f.version, nil
}

func (f *fakeStore) LookupIdempotency(_ context.Context, key string) (string, bool, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	v, ok := f.idem[key]
	return v, ok, nil
}

func (f *fakeStore) SaveIdempotency(_ context.Context, key, nodeUUID string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	if _, ok := f.idem[key]; !ok {
		f.idem[key] = nodeUUID
	}
	return nil
}

func (f *fakeStore) CreateReplacement(_ context.Context, r *Replacement) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	cp := *r
	f.replaces[r.UUID] = &cp
	return nil
}

func (f *fakeStore) GetReplacement(_ context.Context, uuid string) (*Replacement, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	r, ok := f.replaces[uuid]
	if !ok {
		return nil, nil
	}
	cp := *r
	return &cp, nil
}

func (f *fakeStore) UpdateReplacement(_ context.Context, r *Replacement) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	cp := *r
	f.replaces[r.UUID] = &cp
	return nil
}

// helpers for assertions

func (f *fakeStore) countEvents(e Event) int {
	f.mu.Lock()
	defer f.mu.Unlock()
	n := 0
	for _, ev := range f.events {
		if ev.Event == e {
			n++
		}
	}
	return n
}

// firstSeqForStatus returns the seq of the first transition of nodeID into status,
// or -1 if it never happened.
func (f *fakeStore) firstSeqForStatus(nodeID int64, status Status) int {
	f.mu.Lock()
	defer f.mu.Unlock()
	for _, t := range f.transitions {
		if t.NodeID == nodeID && t.Status == status {
			return t.Seq
		}
	}
	return -1
}

// ---------------------------------------------------------------------------
// fakeAdapter / fakeFactory
// ---------------------------------------------------------------------------

type fakeAdapter struct {
	mu           sync.Mutex
	createCalls  int
	destroyCalls int
	createErr    error
	elastic      bool
	nextIP       int
	created      []string // instance IDs created
	destroyed    []string
}

func newFakeAdapter() *fakeAdapter { return &fakeAdapter{elastic: true} }

func (a *fakeAdapter) Kind() string            { return "fake" }
func (a *fakeAdapter) SupportsElasticIP() bool { return a.elastic }

func (a *fakeAdapter) CreateInstance(_ context.Context, in CreateInput) (*Instance, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.createCalls++
	if a.createErr != nil {
		return nil, a.createErr
	}
	a.nextIP++
	id := fmt.Sprintf("inst-%d", a.createCalls)
	a.created = append(a.created, id)
	return &Instance{ID: id, IP: fmt.Sprintf("203.0.113.%d", a.nextIP), Region: in.Region}, nil
}

func (a *fakeAdapter) DestroyInstance(_ context.Context, instanceID string) error {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.destroyCalls++
	a.destroyed = append(a.destroyed, instanceID)
	return nil
}

func (a *fakeAdapter) AttachIP(_ context.Context, _ string) (string, error) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.nextIP++
	return fmt.Sprintf("198.51.100.%d", a.nextIP), nil
}

func (a *fakeAdapter) ListRegions(context.Context) ([]Region, error) {
	return []Region{{ID: "hkg", Country: "HK"}}, nil
}

func (a *fakeAdapter) createCount() int {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.createCalls
}

func (a *fakeAdapter) destroyCount() int {
	a.mu.Lock()
	defer a.mu.Unlock()
	return a.destroyCalls
}

type fakeFactory struct{ a CloudAdapter }

func (f fakeFactory) For(*Provider) (CloudAdapter, error) { return f.a, nil }

// ---------------------------------------------------------------------------
// fakeProber / fakeClock / fakeAlert / fakeTokens
// ---------------------------------------------------------------------------

type fakeProber struct {
	err   error
	calls int
}

func (p *fakeProber) WaitReady(ctx context.Context, _ *Node) error {
	p.calls++
	if p.err != nil {
		return p.err
	}
	return nil
}

// fakeClock records sleeps without actually blocking.
type fakeClock struct {
	now   time.Time
	slept []time.Duration
}

func (c *fakeClock) Now() time.Time { return c.now }
func (c *fakeClock) Sleep(_ context.Context, d time.Duration) error {
	c.slept = append(c.slept, d)
	c.now = c.now.Add(d)
	return nil
}

type fakeAlert struct {
	mu     sync.Mutex
	alerts []Alert
}

func (a *fakeAlert) Fire(_ context.Context, al Alert) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.alerts = append(a.alerts, al)
}

func (a *fakeAlert) count(kind AlertKind) int {
	a.mu.Lock()
	defer a.mu.Unlock()
	n := 0
	for _, al := range a.alerts {
		if al.Kind == kind {
			n++
		}
	}
	return n
}

type fakeTokens struct{ calls int }

func (t *fakeTokens) IssueToken(context.Context, string) (string, error) {
	t.calls++
	return fmt.Sprintf("token-%d", t.calls), nil
}

// ---------------------------------------------------------------------------
// test harness builder
// ---------------------------------------------------------------------------

type harness struct {
	svc     *Service
	store   *fakeStore
	adapter *fakeAdapter
	prober  *fakeProber
	clock   *fakeClock
	alert   *fakeAlert
	tokens  *fakeTokens
}

func newHarness(t interface{ Fatalf(string, ...any) }) *harness {
	store := newFakeStore()
	// One consumable + one premium provider.
	store.addProvider(&Provider{ID: 1, Name: "v", APIKind: "fake", Pool: PoolConsumable, Enabled: true})
	store.addProvider(&Provider{ID: 2, Name: "h", APIKind: "fake", Pool: PoolPremium, Enabled: true})

	adapter := newFakeAdapter()
	prober := &fakeProber{}
	clock := &fakeClock{now: time.Unix(1700000000, 0).UTC()}
	alert := &fakeAlert{}
	tokens := &fakeTokens{}
	renderer, err := NewTemplateRendererFromString("uuid={{.NodeUUID}} token={{.BootstrapToken}}")
	if err != nil {
		t.Fatalf("renderer: %v", err)
	}
	svc, err := NewService(Config{
		Store:           store,
		Adapters:        fakeFactory{adapter},
		Tokens:          tokens,
		Renderer:        renderer,
		Prober:          prober,
		Alert:           alert,
		Clock:           clock,
		ControlPlaneURL: "https://cp.example",
		DrainTimeout:    30 * time.Minute,
		ProbeTimeout:    time.Minute,
	})
	if err != nil {
		t.Fatalf("NewService: %v", err)
	}
	return &harness{svc, store, adapter, prober, clock, alert, tokens}
}

func proSpec() NodeSpec {
	return NodeSpec{Region: "hkg", Tier: TierPro, ProviderID: 2, Plan: "cpx11", RealityPBK: "pbk", RealitySNI: "www.example.com", HY2Port: 443, NameZH: "香港", NameEn: "HK"}
}

func freeSpec() NodeSpec {
	return NodeSpec{Region: "hkg", Tier: TierFree, ProviderID: 1, Plan: "vc2-1c-1gb", RealityPBK: "pbk", RealitySNI: "www.example.com", HY2Port: 443, NameZH: "香港", NameEn: "HK"}
}