Skip to content

Commit

Permalink
fix(MockResponse): added thread-safe to the MockResponse (#155)
Browse files Browse the repository at this point in the history
  • Loading branch information
nomadesoft authored Aug 30, 2024
1 parent 12fca27 commit 913343f
Showing 1 changed file with 30 additions and 0 deletions.
30 changes: 30 additions & 0 deletions mocks.go
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,7 @@ func buildResponseFromMock(mockResponse *MockResponse) *http.Response {
return nil
}

mockResponse.mu.RLock() // Lock for reading
contentTypeHeader := mockResponse.headers["Content-Type"]
var contentType string

Expand All @@ -190,6 +191,8 @@ func buildResponseFromMock(mockResponse *MockResponse) *http.Response {
}
}

mockResponse.mu.RUnlock() // Unlock after reading

res := &http.Response{
Body: ioutil.NopCloser(strings.NewReader(mockResponse.body)),
Header: mockResponse.headers,
Expand All @@ -206,7 +209,9 @@ func buildResponseFromMock(mockResponse *MockResponse) *http.Response {
}

if contentType != "" {
mockResponse.mu.Lock() // Lock for writing
res.Header.Set("Content-Type", contentType)
mockResponse.mu.Unlock() // Unlock after writing
}

return res
Expand Down Expand Up @@ -287,6 +292,7 @@ type MockResponse struct {
body string
statusCode int
fixedDelayMillis int64
mu sync.RWMutex // Add a mutex for thread-safe access
}

// StandaloneMocks for using mocks outside of API tests context
Expand Down Expand Up @@ -634,13 +640,19 @@ func (r *MockResponse) Timeout() *MockResponse {

// Header respond with the given header
func (r *MockResponse) Header(key string, value string) *MockResponse {
r.mu.Lock()
defer r.mu.Unlock()

normalizedKey := textproto.CanonicalMIMEHeaderKey(key)
r.headers[normalizedKey] = append(r.headers[normalizedKey], value)
return r
}

// Headers respond with the given headers
func (r *MockResponse) Headers(headers map[string]string) *MockResponse {
r.mu.Lock()
defer r.mu.Unlock()

for k, v := range headers {
normalizedKey := textproto.CanonicalMIMEHeaderKey(k)
r.headers[normalizedKey] = append(r.headers[normalizedKey], v)
Expand All @@ -650,18 +662,27 @@ func (r *MockResponse) Headers(headers map[string]string) *MockResponse {

// Cookies respond with the given cookies
func (r *MockResponse) Cookies(cookie ...*Cookie) *MockResponse {
r.mu.Lock()
defer r.mu.Unlock()

r.cookies = append(r.cookies, cookie...)
return r
}

// Cookie respond with the given cookie
func (r *MockResponse) Cookie(name, value string) *MockResponse {
r.mu.Lock()
defer r.mu.Unlock()

r.cookies = append(r.cookies, NewCookie(name).Value(value))
return r
}

// Body sets the mock response body
func (r *MockResponse) Body(body string) *MockResponse {
r.mu.Lock()
defer r.mu.Unlock()

r.body = body
return r
}
Expand All @@ -673,6 +694,9 @@ func (r *MockResponse) Bodyf(format string, args ...interface{}) *MockResponse {

// BodyFromFile defines the mock response body from a file
func (r *MockResponse) BodyFromFile(f string) *MockResponse {
r.mu.Lock()
defer r.mu.Unlock()

b, err := ioutil.ReadFile(f)
if err != nil {
panic(err)
Expand All @@ -697,6 +721,9 @@ func (r *MockResponse) JSON(v interface{}) *MockResponse {

// Status respond with the given status
func (r *MockResponse) Status(statusCode int) *MockResponse {
r.mu.Lock()
defer r.mu.Unlock()

r.statusCode = statusCode
return r
}
Expand All @@ -711,6 +738,9 @@ func (r *MockResponse) FixedDelay(delay int64) *MockResponse {

// Times respond the given number of times
func (r *MockResponse) Times(times int) *MockResponse {
r.mu.Lock()
defer r.mu.Unlock()

r.mock.times = times
r.mock.timesSet = true
return r
Expand Down

0 comments on commit 913343f

Please sign in to comment.