Function createChunkedResponseHandler

storage/remote/client_test.go:963–1008

createChunkedResponseHandler creates a mock handler for chunked responses.

(t *testing.T, queries []*prompb.Query)

// createChunkedResponseHandler creates a mock handler for chunked responses.
func createChunkedResponseHandler(t *testing.T, queries []*prompb.Query) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/x-streamed-protobuf; proto=prometheus.ChunkedReadResponse")

		flusher, ok := w.(http.Flusher)
		require.True(t, ok)

		cw := NewChunkedWriter(w, flusher)

		// For each query, simulate multiple chunks
		for queryIndex := range queries {
			chunks := buildTestChunks(t) // Creates 3 chunks with 5 samples each
			for chunkIndex, chunk := range chunks {
				// Create unique labels for each series in each query
				var labels []prompb.Label
				if queryIndex == 0 {
					labels = []prompb.Label{
						{Name: "job", Value: "prometheus"},
						{Name: "instance", Value: fmt.Sprintf("localhost:%d", 9090+chunkIndex)},
					}
				} else {
					labels = []prompb.Label{
						{Name: "job", Value: "node_exporter"},
						{Name: "instance", Value: fmt.Sprintf("localhost:%d", 9100+chunkIndex)},
					}
				}

				cSeries := prompb.ChunkedSeries{
					Labels: labels,
					Chunks: []prompb.Chunk{chunk},
				}

				readResp := prompb.ChunkedReadResponse{
					ChunkedSeries: []*prompb.ChunkedSeries{&cSeries},
					QueryIndex:    int64(queryIndex),
				}

				b, err := proto.Marshal(&readResp)
				require.NoError(t, err)

				_, err = cw.Write(b)
				require.NoError(t, err)
			}
		}
	})
}

// createSampledResponseHandler creates a mock handler for sampled responses.
func createSampledResponseHandler(t *testing.T, queries []*prompb.Query) http.HandlerFunc {