Function TestGetBatchBuffer

modules/frontend/queue/queue_test.go:340–402

(t *testing.T)

}

func TestGetBatchBuffer(t *testing.T) {
	tests := []struct {
		name           string
		queueContents  []Request
		requestedCount int
		expectedCount  int
	}{
		{
			name:           "exactly requested count",
			queueContents:  []Request{&mockRequest{}, &mockRequest{}, &mockRequest{}},
			requestedCount: 3,
			expectedCount:  3,
		},
		{
			name:           "less than requested count",
			queueContents:  []Request{&mockRequest{}, &mockRequest{}},
			requestedCount: 3,
			expectedCount:  2,
		},
		{
			name:           "more than requested count",
			queueContents:  []Request{&mockRequest{}, &mockRequest{}, &mockRequest{}, &mockRequest{}},
			requestedCount: 3,
			expectedCount:  3,
		},
		{
			name:           "less than requested count due to biggest weight",
			queueContents:  []Request{&mockRequest{10}},
			requestedCount: 3,
			expectedCount:  1,
		},
		{
			name:           "empty queue",
			queueContents:  []Request{},
			requestedCount: 3,
			expectedCount:  0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			queue := make(chan Request, len(tt.queueContents))
			for _, req := range tt.queueContents {
				queue <- req
			}

			q := &RequestQueue{
				queueLength: prometheus.NewGaugeVec(prometheus.GaugeOpts{
					Name: "test_len",
				}, []string{"user"}),

				batchWeight: prometheus.NewHistogramVec(prometheus.HistogramOpts{
					Name: "test_weight",
				}, []string{"user"}),
			}

			batchBuffer := make([]Request, tt.requestedCount)
			result := q.getBatchBuffer(batchBuffer, "user", queue)