netbird/management/server/http/testing/testing_tools/tools.go

package testing_tools

import (
	"bytes"
	"context"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"os"
	"strconv"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"

	"github.com/netbirdio/netbird/management/server/account"
	nbpeer "github.com/netbirdio/netbird/management/server/peer"
	"github.com/netbirdio/netbird/management/server/posture"
	"github.com/netbirdio/netbird/management/server/types"
	"github.com/netbirdio/netbird/management/server/util"
)

const (
	TestAccountId = "testAccountId"
	TestPeerId    = "testPeerId"
	TestGroupId   = "testGroupId"
	TestKeyId     = "testKeyId"

	TestUserId         = "testUserId"
	TestAdminId        = "testAdminId"
	TestOwnerId        = "testOwnerId"
	TestServiceUserId  = "testServiceUserId"
	TestServiceAdminId = "testServiceAdminId"
	BlockedUserId      = "blockedUserId"
	OtherUserId        = "otherUserId"
	InvalidToken       = "invalidToken"

	NewKeyName   = "newKey"
	NewGroupId   = "newGroupId"
	ExpiresIn    = 3600
	RevokedKeyId = "revokedKeyId"
	ExpiredKeyId = "expiredKeyId"

	ExistingKeyName = "existingKey"

	OperationCreate = "create"
	OperationUpdate = "update"
	OperationDelete = "delete"
	OperationGetOne = "get_one"
	OperationGetAll = "get_all"
)

var BenchmarkDuration = prometheus.NewGaugeVec(
	prometheus.GaugeOpts{
		Name: "benchmark_duration_ms",
		Help: "Benchmark duration per op in ms",
	},
	[]string{"module", "operation", "test_case", "branch"},
)

type TB interface {
	Cleanup(func())
	Helper()
	Errorf(format string, args ...any)
	Fatalf(format string, args ...any)
	TempDir() string
}

// BenchmarkCase defines a single benchmark test case
type BenchmarkCase struct {
	Peers     int
	Groups    int
	Users     int
	SetupKeys int
}

// PerformanceMetrics holds the performance expectations
type PerformanceMetrics struct {
	MinMsPerOpLocal float64
	MaxMsPerOpLocal float64
	MinMsPerOpCICD  float64
	MaxMsPerOpCICD  float64
}

func BuildRequest(t TB, requestBody []byte, requestType, requestPath, user string) *http.Request {
	t.Helper()

	req := httptest.NewRequest(requestType, requestPath, bytes.NewBuffer(requestBody))
	req.Header.Set("Authorization", "Bearer "+user)

	return req
}

func ReadResponse(t *testing.T, recorder *httptest.ResponseRecorder, expectedStatus int, expectResponse bool) ([]byte, bool) {
	t.Helper()

	res := recorder.Result()
	defer res.Body.Close()

	content, err := io.ReadAll(res.Body)
	if err != nil {
		t.Fatalf("Failed to read response body: %v", err)
	}

	if !expectResponse {
		return nil, false
	}

	if status := recorder.Code; status != expectedStatus {
		t.Fatalf("handler returned wrong status code: got %v want %v, content: %s",
			status, expectedStatus, string(content))
	}

	return content, expectedStatus == http.StatusOK
}

func PopulateTestData(b *testing.B, am account.Manager, peers, groups, users, setupKeys int) {
	b.Helper()

	ctx := context.Background()
	acc, err := am.GetAccount(ctx, TestAccountId)
	if err != nil {
		b.Fatalf("Failed to get account: %v", err)
	}

	// Create peers
	for i := 0; i < peers; i++ {
		peerKey, _ := wgtypes.GeneratePrivateKey()
		peer := &nbpeer.Peer{
			ID:       fmt.Sprintf("oldpeer-%d", i),
			DNSLabel: fmt.Sprintf("oldpeer-%d", i),
			Key:      peerKey.PublicKey().String(),
			IP:       netip.MustParseAddr(fmt.Sprintf("100.64.%d.%d", i/256, i%256)),
			Status:   &nbpeer.PeerStatus{LastSeen: time.Now().UTC(), Connected: true},
			UserID:   TestUserId,
		}
		acc.Peers[peer.ID] = peer
	}

	// Create users
	for i := 0; i < users; i++ {
		user := &types.User{
			Id:        fmt.Sprintf("olduser-%d", i),
			AccountID: acc.Id,
			Role:      types.UserRoleUser,
		}
		acc.Users[user.Id] = user
	}

	for i := 0; i < setupKeys; i++ {
		key := &types.SetupKey{
			Id:         fmt.Sprintf("oldkey-%d", i),
			AccountID:  acc.Id,
			AutoGroups: []string{"someGroupID"},
			UpdatedAt:  time.Now().UTC(),
			ExpiresAt:  util.ToPtr(time.Now().Add(ExpiresIn * time.Second)),
			Name:       NewKeyName + strconv.Itoa(i),
			Type:       "reusable",
			UsageLimit: 0,
		}
		acc.SetupKeys[key.Id] = key
	}

	// Create groups and policies
	acc.Policies = make([]*types.Policy, 0, groups)
	for i := 0; i < groups; i++ {
		groupID := fmt.Sprintf("group-%d", i)
		group := &types.Group{
			ID:   groupID,
			Name: fmt.Sprintf("Group %d", i),
		}
		for j := 0; j < peers/groups; j++ {
			peerIndex := i*(peers/groups) + j
			group.Peers = append(group.Peers, fmt.Sprintf("peer-%d", peerIndex))
		}
		acc.Groups[groupID] = group

		// Create a policy for this group
		policy := &types.Policy{
			ID:      fmt.Sprintf("policy-%d", i),
			Name:    fmt.Sprintf("Policy for Group %d", i),
			Enabled: true,
			Rules: []*types.PolicyRule{
				{
					ID:            fmt.Sprintf("rule-%d", i),
					Name:          fmt.Sprintf("Rule for Group %d", i),
					Enabled:       true,
					Sources:       []string{groupID},
					Destinations:  []string{groupID},
					Bidirectional: true,
					Protocol:      types.PolicyRuleProtocolALL,
					Action:        types.PolicyTrafficActionAccept,
				},
			},
		}
		acc.Policies = append(acc.Policies, policy)
	}

	acc.PostureChecks = []*posture.Checks{
		{
			ID:   "PostureChecksAll",
			Name: "All",
			Checks: posture.ChecksDefinition{
				NBVersionCheck: &posture.NBVersionCheck{
					MinVersion: "0.0.1",
				},
			},
		},
	}

	store := am.GetStore()

	err = store.SaveAccount(context.Background(), acc)
	if err != nil {
		b.Fatalf("Failed to save account: %v", err)
	}

}

func EvaluateAPIBenchmarkResults(b *testing.B, testCase string, duration time.Duration, recorder *httptest.ResponseRecorder, module string, operation string) {
	b.Helper()

	if recorder.Code != http.StatusOK {
		b.Fatalf("Benchmark %s failed: unexpected status code %d", testCase, recorder.Code)
	}

	EvaluateBenchmarkResults(b, testCase, duration, module, operation)

}

func EvaluateBenchmarkResults(b *testing.B, testCase string, duration time.Duration, module string, operation string) {
	b.Helper()

	branch := os.Getenv("GIT_BRANCH")
	if branch == "" && os.Getenv("CI") == "true" {
		b.Fatalf("environment variable GIT_BRANCH is not set")
	}

	msPerOp := float64(duration.Nanoseconds()) / float64(b.N) / 1e6

	gauge := BenchmarkDuration.WithLabelValues(module, operation, testCase, branch)
	gauge.Set(msPerOp)

	b.ReportMetric(msPerOp, "ms/op")
}