1 files changed, 9 insertions, 10 deletions
diff --git a/workhorse/internal/upload/artifacts_store_test.go b/workhorse/internal/upload/artifacts_store_test.go
index 7032313fbde..edba13d38ae 100644
--- a/workhorse/internal/upload/artifacts_store_test.go
+++ b/workhorse/internal/upload/artifacts_store_test.go
@@ -206,14 +206,10 @@ func TestUploadHandlerSendingToExternalStorageAndItReturnsAnError(t *testing.T)
 }
 
 func TestUploadHandlerSendingToExternalStorageAndSupportRequestTimeout(t *testing.T) {
-	putCalledTimes := 0
-
+	shutdown := make(chan struct{})
 	storeServerMux := http.NewServeMux()
 	storeServerMux.HandleFunc("/url/put", func(w http.ResponseWriter, r *http.Request) {
-		putCalledTimes++
-		require.Equal(t, "PUT", r.Method)
-		time.Sleep(10 * time.Second)
-		w.WriteHeader(510)
+		<-shutdown
 	})
 
 	responseProcessor := func(w http.ResponseWriter, r *http.Request) {
@@ -221,13 +217,16 @@ func TestUploadHandlerSendingToExternalStorageAndSupportRequestTimeout(t *testin
 	}
 
 	storeServer := httptest.NewServer(storeServerMux)
-	defer storeServer.Close()
+	defer func() {
+		close(shutdown)
+		storeServer.Close()
+	}()
 
 	authResponse := &api.Response{
 		RemoteObject: api.RemoteObject{
 			StoreURL: storeServer.URL + "/url/put",
 			ID:       "store-id",
-			Timeout:  1,
+			Timeout:  0.001,
 		},
 	}
 
@@ -235,8 +234,8 @@ func TestUploadHandlerSendingToExternalStorageAndSupportRequestTimeout(t *testin
 	defer ts.Close()
 
 	response := testUploadArtifactsFromTestZip(t, ts)
-	require.Equal(t, http.StatusInternalServerError, response.Code)
-	require.Equal(t, 1, putCalledTimes, "upload should be called only once")
+	// HTTP status 504 (gateway timeout) proves that the timeout was enforced
+	require.Equal(t, http.StatusGatewayTimeout, response.Code)
 }
 
 func TestUploadHandlerMultipartUploadSizeLimit(t *testing.T) {