Skip to content

Commit

Permalink
fix(client): Support responses without gzip encoding (#36)
Browse files Browse the repository at this point in the history
  • Loading branch information
mgmacias95 authored Apr 2, 2024
1 parent 8431ff2 commit dbd4c34
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 9 deletions.
18 changes: 12 additions & 6 deletions client.go
Original file line number Diff line number Diff line change
Expand Up @@ -167,14 +167,20 @@ func (cli *Client) parseResponse(resp *http.Response) (*Response, error) {
resp.Request.Method, resp.Request.URL.String())
}

// Prepare gzip reader for uncompressing gzipped JSON response
ungzipper, err := gzip.NewReader(resp.Body)
if err != nil {
return nil, err
var reader io.ReadCloser
if resp.Header.Get("Content-Encoding") == "gzip" {
// Prepare gzip reader for uncompressing gzipped JSON response
ungzipper, err := gzip.NewReader(resp.Body)
if err != nil {
return nil, err
}
defer ungzipper.Close()
reader = ungzipper
} else {
reader = resp.Body
}
defer ungzipper.Close()

if err := json.NewDecoder(ungzipper).Decode(apiresp); err != nil {
if err := json.NewDecoder(reader).Decode(apiresp); err != nil {
return nil, err
}

Expand Down
45 changes: 42 additions & 3 deletions vt_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package vt
import (
"compress/gzip"
"encoding/json"
"errors"
"fmt"
"io/ioutil"
"net/http"
Expand Down Expand Up @@ -30,6 +31,7 @@ type TestServer struct {
expectedMethod string
response interface{}
expectedBody string
status int
expectedHeaders map[string]string
}

Expand All @@ -49,6 +51,11 @@ func (ts *TestServer) SetResponse(r interface{}) *TestServer {
return ts
}

func (ts *TestServer) SetStatusCode(s int) *TestServer {
ts.status = s
return ts
}

func (ts *TestServer) SetExpectedBody(body string) *TestServer {
ts.expectedBody = body
return ts
Expand Down Expand Up @@ -94,9 +101,17 @@ func (ts *TestServer) handler(w http.ResponseWriter, r *http.Request) {
return
}
w.Header().Set("Content-Type", "application/json")
gw := gzip.NewWriter(w)
gw.Write(js)
gw.Close()
if ts.status != 0 {
w.WriteHeader(ts.status)
}
if ts.status != 429 {
w.Header().Set("content-encoding", "gzip")
gw := gzip.NewWriter(w)
gw.Write(js)
gw.Close()
} else {
w.Write(js)
}
}

// This tests GET request with passing in a parameter.
Expand Down Expand Up @@ -429,3 +444,27 @@ func TestRequestHeadersOverrideGlobalHeaders(t *testing.T) {
err := c.PostObject(URL("/collection"), o)
assert.NoError(t, err)
}

func TestGetObjectOutOfQuota(t *testing.T) {
ts := NewTestServer(t).
SetExpectedMethod("GET").
SetStatusCode(429).
SetResponse(map[string]interface{}{
"error": map[string]interface{}{
"code": "QuotaExceededError",
"message": "Quota exceeded",
},
})

defer ts.Close()

SetHost(ts.URL)
c := NewClient("apikey")
_, err := c.GetObject(URL("files/abcabcabcabcabc"))
if err != nil {
var vtErr *Error
if !errors.As(err, &vtErr) && err.(Error).Code != "QuotaExceededError" {
t.Fatalf("Error getting object from VT: %s", err)
}
}
}

0 comments on commit dbd4c34

Please # to comment.