Add x-model header integration test

substratusai · Feb 5, 2024 · c141853 · c141853
1 parent c1beaf3
commit c141853
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/tests/integration/integration_test.go b/tests/integration/integration_test.go
@@ -174,7 +174,7 @@ func TestRetryMiddleware(t *testing.T) {
 		"max retries - fails": {
 			serverCodes:    []int{http.StatusServiceUnavailable, http.StatusServiceUnavailable, http.StatusServiceUnavailable, http.StatusBadGateway},
 			expResultCode:  http.StatusBadGateway,
-			expResultBody:  "{\"error\":\"Bad Gateway\"}\n", // note the linebreak
+			expResultBody:  "502",
 			expBackendHits: 4,
 		},
 		"non retryable error code": {
@@ -196,6 +196,13 @@ func TestRetryMiddleware(t *testing.T) {
 			expResultBody:  "200",
 			expBackendHits: 1,
 		},
+		"503 with model header": {
+			serverCodes:    []int{http.StatusServiceUnavailable, http.StatusServiceUnavailable, http.StatusServiceUnavailable, http.StatusServiceUnavailable},
+			header:         []tuple{{k: "X-Model", v: modelName}},
+			expResultCode:  http.StatusServiceUnavailable,
+			expResultBody:  "503",
+			expBackendHits: 4,
+		},
 	}
 	for name, spec := range specs {
 		t.Run(name, func(t *testing.T) {