internal/resources/traffic_test.go

Rama main · réplica de lectura del repositorio original.

package resources

import (
	"context"
	"strings"
	"sync"
	"testing"
	"time"

	k8sfake "k8s.io/client-go/kubernetes/fake"

	"github.com/sophotechlabs/spinoza/internal/prom"
)

type countingProxy struct {
	mu    sync.Mutex
	calls int
	body  string
}

type blockingProxy struct {
	entered chan struct{}
	release chan struct{}
}

func (b *blockingProxy) Get(ctx context.Context, _ prom.Target, path string, _ map[string]string) ([]byte, error) {
	if path != "api/v1/query" {
		return []byte(`{"status":"success"}`), nil
	}
	select {
	case b.entered <- struct{}{}:
	default:
	}
	select {
	case <-b.release:
		return []byte(labeledFlow), nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (c *countingProxy) Get(ctx context.Context, _ prom.Target, path string, _ map[string]string) ([]byte, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	if path != "api/v1/query" {
		return []byte(`{"status":"success"}`), nil
	}
	c.mu.Lock()
	c.calls++
	c.mu.Unlock()
	return []byte(c.body), nil
}

func (c *countingProxy) count() int {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.calls
}

const labeledFlow = `{"status":"success","data":{"resultType":"vector","result":[
	{"metric":{"source_namespace":"apps","source_workload":"web",
	"destination_namespace":"data","destination_workload":"redis","verdict":"FORWARDED"},
	"value":[1787933018.510,"7"]}
]}}`

func trafficManager(t *testing.T, proxy prom.Proxy, ttl time.Duration) (*Manager, context.CancelFunc) {
	t.Helper()
	ctx, cancel := context.WithCancel(context.Background())
	client := prom.NewClientWithProxy(
		k8sfake.NewClientset(),
		proxy,
		prom.Target{Namespace: "monitoring", Service: "prometheus", Port: "9090"},
	)
	mgr := NewManager(ctx, Deps{
		Dynamic:    newClient(t),
		Clientset:  k8sfake.NewClientset(),
		Prometheus: client,
		Limits:     Limits{IdleGrace: time.Millisecond, TrafficTTL: ttl},
	})
	return mgr, cancel
}

func TestTrafficSupportSaysItIsNotWiredUp(t *testing.T) {
	mgr, cancel := newManager(t, newClient(t))
	defer cancel()

	support := mgr.TrafficSupport(context.Background())

	if support.Available {
		t.Fatal("traffic was offered without prometheus")
	}
	want := "prometheus is not wired up"
	if support.Reason != want {
		t.Fatalf("reason = %q, want %q", support.Reason, want)
	}
}

func TestTrafficGraphSaysItIsNotWiredUp(t *testing.T) {
	mgr, cancel := newManager(t, newClient(t))
	defer cancel()

	graph := mgr.TrafficGraph(context.Background())

	want := "prometheus is not wired up"
	if graph.Error != want {
		t.Fatalf("error = %q, want %q", graph.Error, want)
	}
	if len(graph.Nodes) != 0 {
		t.Fatalf("nodes = %d, want none", len(graph.Nodes))
	}
}

func TestTrafficGraphReadsThroughToPrometheus(t *testing.T) {
	proxy := &countingProxy{body: labeledFlow}
	mgr, cancel := trafficManager(t, proxy, time.Minute)
	defer cancel()

	graph := mgr.TrafficGraph(context.Background())

	if graph.Error != "" {
		t.Fatalf("graph reported %q", graph.Error)
	}
	if len(graph.Edges) != 1 {
		t.Fatalf("edges = %d, want the one prometheus reported", len(graph.Edges))
	}
	if graph.Edges[0].From != "apps/web" {
		t.Fatalf("edge = %+v", graph.Edges[0])
	}
}

func TestASecondWindowSharesTheFirstTrafficProbe(t *testing.T) {
	proxy := &countingProxy{body: labeledFlow}
	mgr, cancel := trafficManager(t, proxy, time.Minute)
	defer cancel()

	first := mgr.TrafficSupport(context.Background())
	after := proxy.count()
	second := mgr.TrafficSupport(context.Background())

	if !first.Available || !second.Available {
		t.Fatalf("support was refused: %+v and %+v", first, second)
	}
	if proxy.count() != after {
		t.Fatalf("prometheus was asked %d times, want the cached answer reused", proxy.count())
	}
}

func TestTheTrafficProbeIsAskedAgainOnceItGoesStale(t *testing.T) {
	proxy := &countingProxy{body: labeledFlow}
	mgr, cancel := trafficManager(t, proxy, time.Nanosecond)
	defer cancel()

	mgr.TrafficSupport(context.Background())
	after := proxy.count()
	mgr.TrafficSupport(context.Background())

	if proxy.count() <= after {
		t.Fatalf("prometheus was asked %d times, want a stale answer refreshed", proxy.count())
	}
}

func TestATrafficProbeOnACancelledContextSaysWhy(t *testing.T) {
	proxy := &countingProxy{body: labeledFlow}
	mgr, cancel := trafficManager(t, proxy, time.Minute)
	defer cancel()
	ctx, stop := context.WithCancel(context.Background())
	stop()

	support := mgr.TrafficSupport(ctx)

	if support.Available {
		t.Fatal("a canceled probe was reported as available")
	}
	if !strings.Contains(support.Reason, "context canceled") {
		t.Fatalf("reason = %q, want the cancellation", support.Reason)
	}
}

func TestAWindowCanStopWaitingForAnotherWindowsTrafficProbe(t *testing.T) {
	proxy := &blockingProxy{entered: make(chan struct{}, 1), release: make(chan struct{})}
	mgr, cancel := trafficManager(t, proxy, time.Minute)
	defer cancel()
	firstDone := make(chan struct{})
	var firstAvailable bool
	go func() {
		firstAvailable = mgr.TrafficSupport(context.Background()).Available
		close(firstDone)
	}()
	<-proxy.entered
	waiting, stop := context.WithCancel(context.Background())
	stop()

	second := mgr.TrafficSupport(waiting)

	if second.Available || !strings.Contains(second.Reason, "context canceled") {
		t.Fatalf("support = %+v, want the canceled waiter released", second)
	}
	close(proxy.release)
	<-firstDone
	if !firstAvailable {
		t.Fatal("canceling the waiter canceled the shared probe")
	}
}

func TestTrafficLimitsFallBackToADefaultTTL(t *testing.T) {
	limits := Limits{}.orDefaults()

	if limits.TrafficTTL != defaultTrafficTTL {
		t.Fatalf("TrafficTTL = %v, want %v", limits.TrafficTTL, defaultTrafficTTL)
	}
}