agent-stack-k8s/internal/controller/limiter/limiter_test.go at d36d4288f17bd1e2dcdec4302ec15b7335d3904b · buildkite/agent-stack-k8s · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
package limiter_test

import (
	"context"
	"errors"
	"testing"

	"github.com/buildkite/agent-stack-k8s/v2/api"
	"github.com/buildkite/agent-stack-k8s/v2/internal/controller/limiter"
	"github.com/buildkite/agent-stack-k8s/v2/internal/controller/model"

	"github.com/google/uuid"
	"go.uber.org/zap/zaptest"
)

func TestLimiter(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	fakeSched := model.NewFakeScheduler(1, nil)
	limiter := limiter.New(ctx, zaptest.NewLogger(t), fakeSched, 1, 1, -1)
	fakeSched.EventHandler = limiter
	fakeSched.Add(50)

	// simulate receiving a bunch of jobs
	var jobs []*api.AgentScheduledJob
	for range 50 {
		jobs = append(jobs, &api.AgentScheduledJob{ID: uuid.New().String()})
	}
	if err := limiter.HandleMany(ctx, jobs); err != nil {
		t.Errorf("limiter.HandleMany(ctx, jobs) = %v", err)
	}
	fakeSched.Wait()

	if got, want := len(fakeSched.Running), 0; got != want {
		t.Errorf("len(fakeSched.Running) = %d, want %d", got, want)
	}
	if got, want := len(fakeSched.Finished), 50; got != want {
		t.Errorf("len(fakeSched.Finished) = %d, want %d", got, want)
	}
	if got, want := fakeSched.Errors, 0; got != want {
		t.Errorf("fakeSched.Errors = %d, want %d", got, want)
	}

	// Wait for the limiter workers to exit (to avoid logging after the test)
	cancel()
	limiter.Wait()
}

func TestLimiter_SchedulerErrors(t *testing.T) {
	t.Parallel()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	fakeSched := model.NewFakeScheduler(0, errors.New("invalid"))
	limiter := limiter.New(ctx, zaptest.NewLogger(t), fakeSched, 1, 1, -1)
	fakeSched.EventHandler = limiter
	fakeSched.Add(50)

	var jobs []*api.AgentScheduledJob
	for range 50 {
		jobs = append(jobs, &api.AgentScheduledJob{ID: uuid.New().String()})
	}
	if err := limiter.HandleMany(ctx, jobs); err != nil {
		t.Errorf("limiter.HandleMany(ctx, jobs) = %v", err)
	}
	fakeSched.Wait()

	if got, want := len(fakeSched.Running), 0; got != want {
		t.Errorf("len(fakeSched.Running) = %d, want %d", got, want)
	}
	if got, want := len(fakeSched.Finished), 0; got != want {
		t.Errorf("len(fakeSched.Finished) = %d, want %d", got, want)
	}
	if got, want := fakeSched.Errors, 50; got != want {
		t.Errorf("fakeSched.Errors = %d, want %d", got, want)
	}

	// Wait for the limiter workers to exit (to avoid logging after the test)
	cancel()
	limiter.Wait()
}