stats-gateway/internal/service/service_test.go

package service

import (
	"context"
	"net"
	"path/filepath"
	"testing"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/status"
	"google.golang.org/grpc/test/bufconn"

	statsv1 "github.com/bergabruh/stats-gateway/gen/stats/v1"
	"github.com/bergabruh/stats-gateway/internal/auth"
	"github.com/bergabruh/stats-gateway/internal/probes"
	"github.com/bergabruh/stats-gateway/internal/storage"
	"github.com/bergabruh/stats-gateway/internal/tunnel"
)

func newTestServer(t *testing.T, st *storage.Store) (statsv1.SpeedStatusClient, func()) {
	t.Helper()
	verifier := auth.NewVerifier([]byte("test-secret"), nil)
	handler := tunnel.NewHandler(tunnel.EgressConfig{Allowed: []string{"127.0.0.1:0"}, RateMbps: 50})

	lis := bufconn.Listen(1024 * 1024)
	srv := grpc.NewServer()
	statsv1.RegisterSpeedStatusServer(srv, New(st, verifier, handler))
	go srv.Serve(lis)

	conn, err := grpc.NewClient("passthrough:///bufnet",
		grpc.WithContextDialer(func(_ context.Context, _ string) (net.Conn, error) { return lis.Dial() }),
		grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		t.Fatal(err)
	}

	cleanup := func() { conn.Close(); srv.Stop() }
	return statsv1.NewSpeedStatusClient(conn), cleanup
}

func TestRecentResults_StreamsAllSamples(t *testing.T) {
	dir := t.TempDir()
	st, _ := storage.Open(filepath.Join(dir, "test.db"))
	defer st.Close()
	for i := 0; i < 3; i++ {
		st.InsertSample(context.Background(), "fl", &probes.Result{
			HTTPLatencyMs: float64(40 + i), Timestamp: time.Now().UTC().Add(time.Duration(i) * time.Second),
		})
	}

	client, cleanup := newTestServer(t, st)
	defer cleanup()

	stream, err := client.RecentResults(context.Background(), &statsv1.Filter{Node: "fl", Limit: 10})
	if err != nil {
		t.Fatal(err)
	}

	count := 0
	for {
		_, err := stream.Recv()
		if err != nil {
			break
		}
		count++
	}
	if count != 3 {
		t.Errorf("expected 3 samples, got %d", count)
	}
}

func TestAggregate_BucketCounts(t *testing.T) {
	dir := t.TempDir()
	st, _ := storage.Open(filepath.Join(dir, "test.db"))
	defer st.Close()

	now := time.Now().UTC()
	// latencies [10, 20, 30, 50, 100, 200] -> buckets [<=25:2, <=50:2, <=100:1, <=250:1, ...]
	latencies := []float64{10, 20, 30, 50, 100, 200}
	for i, ms := range latencies {
		st.InsertSample(context.Background(), "fl", &probes.Result{
			HTTPLatencyMs: ms, Timestamp: now.Add(time.Duration(i) * time.Second),
		})
	}

	client, cleanup := newTestServer(t, st)
	defer cleanup()

	hist, err := client.Aggregate(context.Background(), &statsv1.Range{
		Node:  "fl",
		Since: now.Add(-time.Minute).Format(time.RFC3339),
		Until: now.Add(time.Minute).Format(time.RFC3339),
	})
	if err != nil {
		t.Fatal(err)
	}

	// Find bucket counts by upper_ms
	got := map[float64]int64{}
	for _, b := range hist.Buckets {
		got[b.UpperMs] = b.Count
	}
	want := map[float64]int64{25: 2, 50: 2, 100: 1, 250: 1, 500: 0, 1000: 0, 5000: 0}
	for ms, want := range want {
		if got[ms] != want {
			t.Errorf("bucket <=%.0fms: got %d, want %d", ms, got[ms], want)
		}
	}
}

func TestTunnel_RejectsMissingAuth(t *testing.T) {
	dir := t.TempDir()
	st, _ := storage.Open(filepath.Join(dir, "test.db"))
	defer st.Close()
	client, cleanup := newTestServer(t, st)
	defer cleanup()

	stream, err := client.Tunnel(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	stream.Send(&statsv1.Frame{Kind: statsv1.Frame_HELLO, Target: "1.2.3.4:80"})
	_, err = stream.Recv()
	if err == nil {
		t.Fatal("expected unauthenticated")
	}
	st_, ok := status.FromError(err)
	if !ok || st_.Code() != codes.Unauthenticated {
		t.Errorf("expected Unauthenticated, got %v", err)
	}
}

func TestTunnel_RejectsBadToken(t *testing.T) {
	dir := t.TempDir()
	st, _ := storage.Open(filepath.Join(dir, "test.db"))
	defer st.Close()
	client, cleanup := newTestServer(t, st)
	defer cleanup()

	ctx := metadata.AppendToOutgoingContext(context.Background(), "authorization", "Bearer not.a.real.token")
	stream, err := client.Tunnel(ctx)
	if err != nil {
		t.Fatal(err)
	}
	stream.Send(&statsv1.Frame{Kind: statsv1.Frame_HELLO, Target: "1.2.3.4:80"})
	_, err = stream.Recv()
	if err == nil {
		t.Fatal("expected auth fail")
	}
	st_, ok := status.FromError(err)
	if !ok || st_.Code() != codes.Unauthenticated {
		t.Errorf("expected Unauthenticated, got %v", err)
	}
}