From ebeeb1c8235ec6b9fb1ca334fb08ef304f5bc8e1 Mon Sep 17 00:00:00 2001 From: Patryk Prus Date: Tue, 18 Mar 2025 13:32:17 -0400 Subject: [PATCH] Fix flaky TestReadClient/timeout test Signed-off-by: Patryk Prus --- storage/remote/client_test.go | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/storage/remote/client_test.go b/storage/remote/client_test.go index 7753de04a9..112e96d2b6 100644 --- a/storage/remote/client_test.go +++ b/storage/remote/client_test.go @@ -332,9 +332,9 @@ func TestReadClient(t *testing.T) { }, { name: "timeout", - httpHandler: sampledResponseHTTPHandler(t), - timeout: time.Nanosecond, - expectedErrorContains: "context deadline exceeded: request timed out after 1ns", + httpHandler: delayedResponseHTTPHandler(t, 15*time.Millisecond), + timeout: 5 * time.Millisecond, + expectedErrorContains: "context deadline exceeded: request timed out after 5ms", }, } @@ -448,3 +448,16 @@ func sampledResponseHTTPHandler(t *testing.T) http.HandlerFunc { require.NoError(t, err) } } + +func delayedResponseHTTPHandler(t *testing.T, delay time.Duration) http.HandlerFunc { + return func(w http.ResponseWriter, _ *http.Request) { + time.Sleep(delay) + + w.Header().Set("Content-Type", "application/x-protobuf") + b, err := proto.Marshal(&prompb.ReadResponse{}) + require.NoError(t, err) + + _, err = w.Write(snappy.Encode(nil, b)) + require.NoError(t, err) + } +}