Use HTTP RoundTripper in unit tests for code involves Azure SDK

Azure SDK 里面的 client creation 的时候，允许给入一个 client option，which can be a HTTP transport (which holds an interface type: http.RoundTripper).

这样我们就可以 mock the HTTP traffic → test the code which talks to Azure. 这样我们就可以很大程度上完成 Azure SDK involved E2E test，而且不需要我们为 Azure SDK 写 mock.

Mock HTTP handler

// The requests and the mock responses
type SimpleDeclarativeMap map[DeclarativeRequestKey]DeclarativeResponse
 
// SimpleDeclarativeHandler is an implementation of http.Handler that returns
// stubbed responses defined in a map of request/response pairs. Requests are
// matched verbatim against method/pathname/body, and the corresponding canned
// response is returned, or a 500 if no match was found.
type SimpleDeclarativeHandler struct {
  Map  SimpleDeclarativeMap
  Seen map[DeclarativeRequestKey]struct{}
}
 
// This is implementating the http.Handler interface.
// Note that this is NOT the 'RoundTripper' interface.
func (h SimpleDeclarativeHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
  resp, err := func() (DeclarativeResponse, error) {
    key, err := newDeclarativeRequestKey(req)
    // handle err
    resp, ok := h.Map[key]
    // if no mock response found, return error, otherwise, return the response
    // and mark the 'Seen' map.
  }()
 
  // If the functor above returns an error, we should write HTTP internal error
  // to the status header to mark the 'HTTP error' for the client.
  if err != nil {
    w.WriteHeader(http.StatusInternalServerError); fmt.Fprint(w, err.Error())
    return
  }
  // Write the header
  for k, v := range resp.Header { w.Header().Set(k, v) }
  // Write status header
  w.WriteHeader(resp.Status)
  // Write body
  fmt.Fprint(w, resp.Body)
  // we are done.
}
 
// Compare the 'Seen' and 'Map' to check all keys in Map should be in Seen.
func (h SimpleDeclarativeHandler) AssertNoUnmatchedRequests(t *testing.T) {
  for k := range h.Map {
    assert.Containsf(t, h.Seen, k, "Unmatched request: %s", k)
  }
}
 
// -----------------------------------------------
// Internal Details for the request, response, etc.
// -----------------------------------------------
 
// This is also used as the 'key' to determine whether a specific request
// has been made to Azure.
type DeclarativeRequestKey struct {
  Method, Pathname, Body string
}
 
func newDeclarativeRequestKey(req *http.Request) (DeclarativeRequestKey, error) {
  // Body: io.ReadAll(req.Body)
  // Path: &url.URL{Path: request.URL.Path, RawQuery: request.URL.RawQuery} Note: host, scheme are not important
  // Method: req.Method
}
 
func (k DeclarativeRequestKey) String() string { /* 'k.Method'.'k.Pathname'.'k.Body' */ }
 
// This is the mock response
type DeclarativeResponse struct {
  Status int
  Body   string
  Header map[string]string
}
 
// The handler ctor, prevent 'segfault' due to uninitialized map.
func NewSimpleDeclarativeHandler(m SimpleDeclarativeMap) SimpleDeclarativeHandler {
  return SimpleDeclarativeHandler{ Map: m, Seen: ..., }
}

NOTE

This SimpleDeclarativeHandler implements the http.Handler interface. Not that http.RoundTripper interface.

How to use this SimpleDeclarativeHandler

We need to somehow create a http.RoundTripper from http.Handler.

// tt should be the 'testcase struct'.
mockHandler := NewSimpleDeclarativeHandler(tc.RequestMap)
apiVerStrippedHandler := IgnoreAzureAPIVersionMiddleeware(mockHandler)
roundTripper := NewTestRoundTripper(apiVerStrippedHandler)
azureClientOpts := &arm.ClientOptions{
  ClientOptions: policy.ClientOptions{
    Transport: mock.AzureTransportAdapter{
      Transport: roundTripper,
    },
  },
}

Here, we have several layers or wrapping:

handler is wrapped with IgnoreAzureAPIVersionMiddleeware

• This middleware is a thing “taking a http.Handler and returning a new http.Handler”.
  • This middleware just “deletes” the api-version query in the URL, this is to make the E2E test more robust.

func IgnoreAzureAPIVersionMiddleware(h http.Handler) http.Handler {
  return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
    req = req.Clone(req.Context())
    query, err := url.ParseQuery(req.URL.RawQuery)
    if err != nil {
      w.WriteHeader(http.StatusInternalServerError)
      fmt.Fprintf(w, "could not parse query: %v", err)
      return
    }
    query.Del("api-version") // avoid having to update test when SDK is updated
    req.URL.RawQuery = query.Encode()
    h.ServeHTTP(w, req)
  })
}

Different types of ‘Transporter’

Azure SDK defines a ‘Transporter’ interface:

type Transporter interface {
  // Do sends the HTTP request and returns the HTTP response or error.
  Do(req *http.Request) (*http.Response, error)
}

This is not the typical ‘http.RoundTripper’ interface which is usually used in as the ‘Transport’ field in standard http.Client.

This interface is actually implemented by http.Client.

So actually we can use http.Client as the Transporter for Azure client option.

Cast HTTP Handler direclty to a HTTP RoundTripper

• The ‘api-version’ stripped handler is used to create the roundTripper, so what is the simplest way to create a http.RoundTripper?

type testRoundTripper struct {
  h http.Handler
}
 
func (t testRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
  req = req.Clone(req.Context())
  // req.URL.Scheme and req.URL.Host are important.
  // But we need to set them. The Azure SDK does verify the 'Request' sent back
  // to the client in the Response and verify the URL is a valid one.
  // - I guess the Azure SDK will formulate the next request URL based on the
  //   previous response for some operations, like op poller.
  req.URL.Scheme = "http"
  req.URL.Host = "localhost:8080"
  rw := httptest.NewRecorder()
  r.handler.ServeHTTP(rw, req)
  resp := rw.Result()
  // We need to assign the request back to the 'Respponse' object to:
  // SIMULATE the behavior of golang HTTP client.
  // If you miss this, you will see panic due to nil pointer in Azure SDK.
  resp.Request = req // This is super important!
  return resp
}
 
func NewTestRoundTripper(h http.Handler) http.RoundTripper {
  return testRoundTripper{h}
}

WARNING

We have to set the ‘resp.Request’ because this the behavior of Golang HTTP client.

Azure Golang SDK of course uses the Golang HTTP client and it does use the Request object in the Response when creating the Async Poller:

• Code of nil pointer dereferencing,
• OpenSource Code of nil pointer dereferencing

So where is the ‘Request’ set back to the ‘Response’ object in Golang’s standard http stack?

• Code where the assignment happens
  • Used in *Transport
    • Used in httptest.Server. and here So that the ‘test’ server will work as expected.
    • Used as the DefaltTransport
      • Note that this is used by the DefaultClient in Golang. (So, again, this is a client side behavior.)

Better: use httptest.NewServer to emulate an actual HTTP server in a better way

Option 1: Wrap and wrap and wrap

• httptest.NewServer will establish a real server on a random tcp port.
• It also provies a ‘client’ to use for testing.
• It allows an easy integration with a custom http.Handler for testing.

type MockServer struct{
  *httptest.Server // so that we 'inherit' all the methods of httptest.Server. Type alias won't work for this.
}
 
// We just want to let this new server be able to act as a round tripper.
// But under the hood, it is a server and it just delegates the request to the
// handler with the client provided by httptest.Server.
func (s MockServer) RoundTrip(req *http.Request) (*http.Response, error) {
  req = req.Clone(req.Context())
  req.URL.Scheme = "http" // do e need this? 
  req.URL.Host = m.Listener.Addr().String() // do we need this?
  return s.Client().Do(req)
}
 
func NewMockServer(h http.Handler) MockServer {
  return MockServer{httptest.NewServer(h)}
}

Then in the test side:

// tt should be the 'testcase struct'.
mockHandler := NewSimpleDeclarativeHandler(tc.RequestMap)
apiVerStrippedHandler := IgnoreAzureAPIVersionMiddleeware(mockHandler)
s := NewMockServer(apiVerStrippedHandler) // this is new
azureClientOpts := &arm.ClientOptions{
  ClientOptions: policy.ClientOptions{
    Transport: AzureTransportAdapter{
      Transport: s, // used as the transport for the client option.
    },
  },
}

See, we still need to use the AzureTransportAdapter to wrap the MockServer to make it work with the Azure SDK.

type AzureTransportAdapter struct {
  Transport http.RoundTripper
}
 
func (a AzureTransportAdapter) Do(req *http.Request) (*http.Response, error) {
  return a.Transport.RoundTrip(req)
}

Option 2 Probably better: just wrap the Client provided by httptest.Server

// Just define a new Client struct which will replace the Schema and Host so the
// request goes to the test server.
type newTestClient struct {
  *http.Client
  addr string
}
 
// Replace the Do method to set the Schema and Host to redirect the request
// to the test server started with httptest.NewServer.
func (c newTestClient) Do(req *http.Request) (*http.Response, error) {
  req = req.Clone(req.Context())
  req.URL.Scheme = "http"
  req.URL.Host = c.addr
  return c.Client.Do(req)
}
 
// A helper function to create the 
func newHttpTestServerClient(s *http.Server) *newTestClient {
  origClient := s.Client()
  addr := s.Listener.Addr().String()
  c := &newTestClient{
    Client: origClient,
    addr:   addr,
  }
  return c
}

So in the test side:

// tt should be the 'testcase struct'.
mockHandler := NewSimpleDeclarativeHandler(tc.RequestMap)
apiVerStrippedHandler := IgnoreAzureAPIVersionMiddleeware(mockHandler)
s := httptest.NewServer(apiVerStrippedHandler)
defer s.Close()
azureClientOpts := &arm.ClientOptions{
  ClientOptions: policy.ClientOptions{
    Transport: newHttpTestServerClient(s),
  },
}

How to use this in the test

func TestLaunchHost(t *testing.T) {
  tests := []struct {
    name                               string
    params                             Params
    requestMap                         httptest.SimpleDeclarativeMap
    expectLaunchHostError              bool
    launchHostErrorContains            string
    expectedVmName                     string
    expectedComputerName               string
    expectedLaunchInfo                 LaunchInfo
    expectedPollUntilDoneError         bool
    expectedPollUntilDoneErrorContains string
  }{
    {
      name: "successful launch",
      params: Params{...},
      requestMap: httptest.SimpleDeclarativeMap{
        // 'Get VMSS' Expected Request and Mock Response
        {
          Method:   http.MethodGet, Pathname: ..., /* the pathname expected in the http handler. */
        }: {
          Status: http.StatusOK, Body: ... , /* the mock response body. */
        },
        // 'Create VM' Expected Request and Mock Response
        {
          Method:   http.MethodPut,
          Pathname: ..., /* the pathname expected in the http handler. */
          Body:     ..., /* the expected request body. */
        }: {
          Status: http.StatusCreated,
          Body:   ..., /* the mock response body. */
          Header: ..., /* the mock response header. */
        },
        // 'Poll' Expected Request and Mock Response
        {
          Method:   http.MethodGet,
          Pathname: ..., /* the pathname expected in the http handler. */
        }: {
          Status: http.StatusOK,
          Body:   ..., /* the mock response body. */
        },
        // 'Get VM' Expected Request and Mock Response
        {
          Method:   http.MethodGet,
          Pathname: ..., /* the pathname expected in the http handler. */
        }: {
          Status: http.StatusOK,
          Body:   ..., /* the mock response body. */
        },
      },
      expectXXX: ...,
    },
    // ... more test cases
  }
  for _, tt := range tests {
    t.Run(tt.name, func(t *testing.T) {
      // Create HTTP test server with mock responses
      mockHandler := NewSimpleDeclarativeHandler(tc.RequestMap)
      apiVerStrippedHandler := IgnoreAzureAPIVersionMiddleeware(mockHandler)
      s := httptest.NewServer(apiVerStrippedHandler)
      defer func() {
        // Verify all expected HTTP requests were made
        handler.AssertNoUnmatchedRequests(t)
      }()
      defer s.Close()
 
      // Create HostLauncher with mock transport and token
      launcher := &HostLauncher{...
        // Use http round tripper to capture HTTP requests and mock
        // HTTP responses.
        azureClientOptions: &arm.ClientOptions{
          ClientOptions: policy.ClientOptions{
            Transport: newHttpTestServerClient(s),
          },
        },
        // More settings...
      }
 
      // Test Launch host
      op, err := launcher.LaunchHost(ctx, tt.params)
      // .... Handling the result and error ...
    })
  }
}

Interesting Bits

• On the Client side, for Golang, the ‘original request’ is set back in the ‘Response’
  • On the Sever side, there is no such behavior.
  • Both the Client side and the Server side uses the same http.Response type (this is not an interface).
• The httptest.NewRecorder() is the easiest way to create a http.ResponseWriter to be used with http.Handler (i.e. ServeHTTP() method).
• The Client() return of httptest.NewServer() does not override the Schema and the Host fields automatically for an incoming request.
  • If we want the request to hit the httptest.NewServer(), we still need to manually set the schema and the host.
• Azure SDK use
  • Azure-AsyncOperation header to get the URL to poll the long running operation status;
  • Location header to determine the ‘underlying resource’ processed by the long running operation. Once the long running operation is done, the poller (in the Azure SDK) will GET the Location URL to get the ‘resource’ (e.g. a VM).