github.com/google/martian/v3@v3.3.3/martiantest/martiantest.go (about)

     1  // Copyright 2015 Google Inc. All rights reserved.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package martiantest provides helper utilities for testing
    16  // modifiers.
    17  package martiantest
    18  
    19  import (
    20  	"net/http"
    21  	"sync/atomic"
    22  )
    23  
    24  // Modifier keeps track of the number of requests and responses it has modified
    25  // and can be configured to return errors or run custom functions.
    26  type Modifier struct {
    27  	reqcount int32 // atomic
    28  	rescount int32 // atomic
    29  	reqerr   error
    30  	reserr   error
    31  	reqfunc  func(*http.Request)
    32  	resfunc  func(*http.Response)
    33  }
    34  
    35  // NewModifier returns a new test modifier.
    36  func NewModifier() *Modifier {
    37  	return &Modifier{}
    38  }
    39  
    40  // RequestCount returns the number of requests modified.
    41  func (m *Modifier) RequestCount() int32 {
    42  	return atomic.LoadInt32(&m.reqcount)
    43  }
    44  
    45  // ResponseCount returns the number of responses modified.
    46  func (m *Modifier) ResponseCount() int32 {
    47  	return atomic.LoadInt32(&m.rescount)
    48  }
    49  
    50  // RequestModified returns whether a request has been modified.
    51  func (m *Modifier) RequestModified() bool {
    52  	return m.RequestCount() != 0
    53  }
    54  
    55  // ResponseModified returns whether a response has been modified.
    56  func (m *Modifier) ResponseModified() bool {
    57  	return m.ResponseCount() != 0
    58  }
    59  
    60  // RequestError overrides the error returned by ModifyRequest.
    61  func (m *Modifier) RequestError(err error) {
    62  	m.reqerr = err
    63  }
    64  
    65  // ResponseError overrides the error returned by ModifyResponse.
    66  func (m *Modifier) ResponseError(err error) {
    67  	m.reserr = err
    68  }
    69  
    70  // RequestFunc is a function to run during ModifyRequest.
    71  func (m *Modifier) RequestFunc(reqfunc func(req *http.Request)) {
    72  	m.reqfunc = reqfunc
    73  }
    74  
    75  // ResponseFunc is a function to run during ModifyResponse.
    76  func (m *Modifier) ResponseFunc(resfunc func(res *http.Response)) {
    77  	m.resfunc = resfunc
    78  }
    79  
    80  // ModifyRequest increases the count of requests seen and runs reqfunc if configured.
    81  func (m *Modifier) ModifyRequest(req *http.Request) error {
    82  	atomic.AddInt32(&m.reqcount, 1)
    83  
    84  	if m.reqfunc != nil {
    85  		m.reqfunc(req)
    86  	}
    87  
    88  	return m.reqerr
    89  }
    90  
    91  // ModifyResponse increases the count of responses seen and runs resfunc if configured.
    92  func (m *Modifier) ModifyResponse(res *http.Response) error {
    93  	atomic.AddInt32(&m.rescount, 1)
    94  
    95  	if m.resfunc != nil {
    96  		m.resfunc(res)
    97  	}
    98  
    99  	return m.reserr
   100  }
   101  
   102  // Reset resets the request and response counts, the custom
   103  // functions, and the modifier errors.
   104  func (m *Modifier) Reset() {
   105  	atomic.StoreInt32(&m.reqcount, 0)
   106  	atomic.StoreInt32(&m.rescount, 0)
   107  
   108  	m.reqfunc = nil
   109  	m.resfunc = nil
   110  
   111  	m.reqerr = nil
   112  	m.reserr = nil
   113  }
   114  
   115  // Matcher is a stubbed matcher used in tests.
   116  type Matcher struct {
   117  	resval bool
   118  	reqval bool
   119  }
   120  
   121  // NewMatcher returns a pointer to martiantest.Matcher with the return values
   122  // for MatchRequest and MatchResponse intiailized to true.
   123  func NewMatcher() *Matcher {
   124  	return &Matcher{resval: true, reqval: true}
   125  }
   126  
   127  // ResponseEvaluatesTo sets the value returned by MatchResponse.
   128  func (tm *Matcher) ResponseEvaluatesTo(value bool) {
   129  	tm.resval = value
   130  }
   131  
   132  // RequestEvaluatesTo sets the value returned by MatchRequest.
   133  func (tm *Matcher) RequestEvaluatesTo(value bool) {
   134  	tm.reqval = value
   135  }
   136  
   137  // MatchRequest returns the stubbed value in tm.reqval.
   138  func (tm *Matcher) MatchRequest(*http.Request) bool {
   139  	return tm.reqval
   140  }
   141  
   142  // MatchResponse returns the stubbed value in tm.resval.
   143  func (tm *Matcher) MatchResponse(*http.Response) bool {
   144  	return tm.resval
   145  }