Function TestAsyncSharders

modules/frontend/pipeline/async_sharding_test.go:15–95

(t *testing.T)

)

func TestAsyncSharders(t *testing.T) {
	expectedRequestCount := 10000

	tcs := []struct {
		name       string
		responseFn func(next AsyncRoundTripper[combiner.PipelineResponse]) *asyncResponse
	}{
		{
			name: "AsyncSharder",
			responseFn: func(next AsyncRoundTripper[combiner.PipelineResponse]) *asyncResponse {
				return NewAsyncSharderFunc(context.Background(), 10, expectedRequestCount, func(i int) Request {
					if i >= expectedRequestCount {
						return nil
					}
					return NewHTTPRequest(&http.Request{})
				}, next).(*asyncResponse)
			},
		},
		{
			name: "AsyncSharder - no limit",
			responseFn: func(next AsyncRoundTripper[combiner.PipelineResponse]) *asyncResponse {
				return NewAsyncSharderFunc(context.Background(), 0, expectedRequestCount, func(i int) Request {
					if i >= expectedRequestCount {
						return nil
					}
					return NewHTTPRequest(&http.Request{})
				}, next).(*asyncResponse)
			},
		},
		{
			name: "AsyncSharderLimitedGoroutines",
			responseFn: func(next AsyncRoundTripper[combiner.PipelineResponse]) *asyncResponse {
				reqChan := make(chan Request)
				go func() {
					for i := 0; i < expectedRequestCount; i++ {
						reqChan <- NewHTTPRequest(&http.Request{})
					}
					close(reqChan)
				}()

				return NewAsyncSharderChan(context.Background(), 10, reqChan, nil, next).(*asyncResponse)
			},
		},
	}

	for _, tc := range tcs {
		t.Run(tc.name, func(t *testing.T) {
			next := AsyncRoundTripperFunc[combiner.PipelineResponse](func(_ Request) (Responses[combiner.PipelineResponse], error) {
				// return a generic 200
				return NewHTTPToAsyncResponse(&http.Response{
					Body:       io.NopCloser(strings.NewReader("")),
					StatusCode: 200,
				}), nil
			})

			sharderResp := tc.responseFn(next)

			// drain and count the requests