support multiple expectedStatusCodes

vtopc · Aug 1, 2023 · a31d27a · a31d27a
1 parent 78e13dd
commit a31d27a
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 50 deletions.
diff --git a/client.go b/client.go
@@ -52,11 +52,6 @@ func (c *Client) Do(req *http.Request, v interface{}, expectedStatusCodes ...int
 		return errors.New("empty request")
 	}
 
-	// TODO: add support for the multiple status codes
-	if len(expectedStatusCodes) > 1 {
-		return errors.New("support for multiple status codes is not implemented")
-	}
-
 	// Set defaults:
 	if len(expectedStatusCodes) == 0 {
 		expectedStatusCodes = []int{http.StatusOK}
@@ -78,7 +73,15 @@ func (c *Client) do(req *http.Request, v interface{}, expectedStatusCodes ...int
 
 	defer resp.Body.Close()
 
-	if resp.StatusCode != expectedStatusCodes[0] {
+	var expected bool
+	for _, code := range expectedStatusCodes {
+		if resp.StatusCode == code {
+			expected = true
+			break
+		}
+	}
+
+	if !expected {
 		// Non expected status code.
 
 		buf := &strings.Builder{}

diff --git a/client_test.go b/client_test.go
@@ -17,63 +17,73 @@ type Struct struct {
 
 func TestClientDo(t *testing.T) {
 	tests := map[string]struct {
-		method             string
-		urlPostfix         string
-		statusCode         int
-		expectedStatusCode int
-		respBody           []byte
-		v                  interface{}
-		want               interface{}
-		wantErr            error
-		wantWrappedErr     error
+		method              string
+		urlPostfix          string
+		statusCode          int
+		expectedStatusCodes []int
+		respBody            []byte
+		v                   interface{}
+		want                interface{}
+		wantErr             error
+		wantWrappedErr      error
 	}{
 		"positive_get": {
-			method:             http.MethodGet,
-			urlPostfix:         "/health",
-			statusCode:         http.StatusOK,
-			expectedStatusCode: http.StatusOK,
-			respBody:           []byte(`{"status":"ok"}`),
-			v:                  &Struct{},
-			want:               &Struct{Status: "ok"},
+			method:              http.MethodGet,
+			urlPostfix:          "/health",
+			statusCode:          http.StatusOK,
+			expectedStatusCodes: []int{http.StatusOK},
+			respBody:            []byte(`{"status":"ok"}`),
+			v:                   &Struct{},
+			want:                &Struct{Status: "ok"},
+		},
+
+		"positive_multiple_statuses": {
+			method:              http.MethodPost,
+			urlPostfix:          "/health",
+			statusCode:          http.StatusAccepted,
+			expectedStatusCodes: []int{http.StatusOK, http.StatusAccepted},
+			respBody:            []byte(`{"status":"ok"}`),
+			v:                   &Struct{},
+			want:                &Struct{Status: "ok"},
 		},
 
 		"positive_but_wrong_payload": {
-			method:             http.MethodGet,
-			urlPostfix:         "/health",
-			statusCode:         http.StatusOK,
-			expectedStatusCode: http.StatusOK,
-			respBody:           []byte(`{"error":"some error"}`),
-			v:                  &Struct{},
-			want:               &Struct{}, // zero values
+			method:              http.MethodGet,
+			urlPostfix:          "/health",
+			statusCode:          http.StatusOK,
+			expectedStatusCodes: []int{http.StatusOK},
+			respBody:            []byte(`{"error":"some error"}`),
+			v:                   &Struct{},
+			want:                &Struct{}, // zero values
 		},
 
 		"positive_get_empty_body": {
-			method:             http.MethodGet,
-			urlPostfix:         "/health",
-			statusCode:         http.StatusNoContent,
-			expectedStatusCode: http.StatusNoContent,
-			respBody:           nil,
-			v:                  nil,
-			want:               nil,
+			method:              http.MethodGet,
+			urlPostfix:          "/health",
+			statusCode:          http.StatusNoContent,
+			expectedStatusCodes: []int{http.StatusNoContent},
+			respBody:            nil,
+			v:                   nil,
+			want:                nil,
 		},
 
 		"negative_wrong_status_code": {
-			method:             http.MethodGet,
-			urlPostfix:         "/health",
-			statusCode:         http.StatusInternalServerError,
-			expectedStatusCode: http.StatusOK,
-			respBody:           []byte(`{"error":"some error"}`),
-			wantWrappedErr:     errors.New("wrong status code (500 not in [200]): {\"error\":\"some error\"}"),
+			method:              http.MethodGet,
+			urlPostfix:          "/health",
+			statusCode:          http.StatusInternalServerError,
+			expectedStatusCodes: []int{http.StatusOK},
+			respBody:            []byte(`{"error":"some error"}`),
+			wantWrappedErr:      errors.New("wrong status code (500 not in [200]): {\"error\":\"some error\"}"),
 		},
 
 		"negative_not_a_pointer": {
-			method:             http.MethodGet,
-			urlPostfix:         "/health",
-			statusCode:         http.StatusOK,
-			expectedStatusCode: http.StatusOK,
-			respBody:           []byte(`{"status":"ok"}`),
-			v:                  Struct{},
-			wantErr:            errors.New("json: Unmarshal(non-pointer rest.Struct)"),
+			method:              http.MethodGet,
+			urlPostfix:          "/health",
+			statusCode:          http.StatusOK,
+			expectedStatusCodes: []int{http.StatusOK},
+			respBody:            []byte(`{"status":"ok"}`),
+			v:                   Struct{},
+			wantErr:             errors.New("json: Unmarshal(non-pointer rest.Struct)"),
 		},
 
 		// TODO: add more test cases
@@ -104,7 +114,7 @@ func TestClientDo(t *testing.T) {
 			require.NoError(t, err)
 
 			// test:
-			err = c.Do(req, tt.v, tt.expectedStatusCode)
+			err = c.Do(req, tt.v, tt.expectedStatusCodes...)
 			if tt.wantErr != nil {
 				require.EqualError(t, err, tt.wantErr.Error())
 				return