github.com/m3db/m3@v1.5.0/src/query/api/v1/middleware/middleware_test.go (about)

     1  // Copyright (c) 2021 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package middleware
    22  
    23  import (
    24  	"compress/gzip"
    25  	"fmt"
    26  	"io/ioutil"
    27  	"net/http"
    28  	"net/http/httptest"
    29  	"testing"
    30  
    31  	"github.com/gorilla/mux"
    32  	"github.com/opentracing/opentracing-go/mocktracer"
    33  	"github.com/stretchr/testify/assert"
    34  	"github.com/stretchr/testify/require"
    35  	"go.uber.org/zap"
    36  	"go.uber.org/zap/zapcore"
    37  	"go.uber.org/zap/zaptest/observer"
    38  
    39  	"github.com/m3db/m3/src/query/util/logging"
    40  	"github.com/m3db/m3/src/x/instrument"
    41  )
    42  
    43  func TestTracing(t *testing.T) {
    44  	core, recorded := observer.New(zapcore.InfoLevel)
    45  	iOpts := instrument.NewOptions().SetLogger(zap.New(core))
    46  	mtr := mocktracer.New()
    47  	r := mux.NewRouter()
    48  	r.Use(Tracing(mtr, iOpts))
    49  	r.HandleFunc(testRoute, func(w http.ResponseWriter, r *http.Request) {
    50  		logging.WithContext(r.Context(), iOpts).Info("test")
    51  	})
    52  
    53  	req := httptest.NewRequest("GET", testRoute, nil)
    54  	res := httptest.NewRecorder()
    55  	r.ServeHTTP(res, req)
    56  
    57  	spans := mtr.FinishedSpans()
    58  	require.Len(t, spans, 1)
    59  	require.Equal(t, fmt.Sprintf("GET %s", testRoute), spans[0].OperationName)
    60  	require.Len(t, recorded.All(), 1)
    61  	entry := recorded.All()[0]
    62  	require.Equal(t, "test", entry.Message)
    63  	fields := entry.ContextMap()
    64  	require.Len(t, fields, 2)
    65  	require.NotEqual(t, "", fields["trace_id"])
    66  	require.NotEqual(t, "", fields["span_id"])
    67  }
    68  
    69  func TestCompression(t *testing.T) {
    70  	router := mux.NewRouter()
    71  	setupTestRouteRouter(router)
    72  
    73  	router.Use(Compression())
    74  
    75  	req := httptest.NewRequest("GET", testRoute, nil)
    76  	req.Header.Add("Accept-Encoding", "gzip")
    77  	res := httptest.NewRecorder()
    78  	router.ServeHTTP(res, req)
    79  
    80  	enc, found := res.Header()["Content-Encoding"]
    81  	require.True(t, found)
    82  	require.Equal(t, 1, len(enc))
    83  	assert.Equal(t, "gzip", enc[0])
    84  
    85  	cr, err := gzip.NewReader(res.Body)
    86  	require.NoError(t, err)
    87  	body, err := ioutil.ReadAll(cr)
    88  	require.NoError(t, err)
    89  	assert.Equal(t, "hello!", string(body))
    90  }
    91  
    92  func TestCors(t *testing.T) {
    93  	router := mux.NewRouter()
    94  	setupTestRouteRouter(router)
    95  
    96  	router.Use(Cors())
    97  
    98  	req := httptest.NewRequest("GET", testRoute, nil)
    99  	res := httptest.NewRecorder()
   100  	router.ServeHTTP(res, req)
   101  
   102  	assert.Equal(t, "hello!", res.Body.String())
   103  	assert.Equal(t, "*", res.Header().Get("Access-Control-Allow-Origin"))
   104  }
   105  
   106  const testRoute = "/foobar"
   107  
   108  func setupTestRouteRouter(r *mux.Router) {
   109  	r.HandleFunc(testRoute, func(writer http.ResponseWriter, r *http.Request) {
   110  		writer.WriteHeader(http.StatusOK)
   111  		_, _ = writer.Write([]byte("hello!"))
   112  	})
   113  }