github.com/m3db/m3@v1.5.0/src/query/storage/promremote/promremotetest/test_server.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 promremotetest provides test utilities.
    22  package promremotetest
    23  
    24  import (
    25  	"fmt"
    26  	"net/http"
    27  	"net/http/httptest"
    28  	"sync"
    29  	"testing"
    30  
    31  	"github.com/prometheus/prometheus/prompb"
    32  	"github.com/prometheus/prometheus/storage/remote"
    33  	"github.com/stretchr/testify/assert"
    34  )
    35  
    36  // TestPromServer is a fake http server handling prometheus remote write. Intended for test usage.
    37  type TestPromServer struct {
    38  	mu               sync.Mutex
    39  	lastWriteRequest *prompb.WriteRequest
    40  	respErr          *respErr
    41  	t                *testing.T
    42  	svr              *httptest.Server
    43  }
    44  
    45  type respErr struct {
    46  	error  string
    47  	status int
    48  }
    49  
    50  // NewServer creates new instance of a fake server.
    51  func NewServer(t *testing.T) *TestPromServer {
    52  	testPromServer := &TestPromServer{t: t}
    53  
    54  	mux := http.NewServeMux()
    55  	mux.HandleFunc("/write", testPromServer.handleWrite)
    56  
    57  	testPromServer.svr = httptest.NewServer(mux)
    58  
    59  	return testPromServer
    60  }
    61  
    62  func (s *TestPromServer) handleWrite(w http.ResponseWriter, r *http.Request) {
    63  	s.mu.Lock()
    64  	defer s.mu.Unlock()
    65  	assert.Equal(s.t, r.Header.Get("content-encoding"), "snappy")
    66  	assert.Equal(s.t, r.Header.Get("content-type"), "application/x-protobuf")
    67  
    68  	req, err := remote.DecodeWriteRequest(r.Body)
    69  	if err != nil {
    70  		http.Error(w, err.Error(), http.StatusBadRequest)
    71  		return
    72  	}
    73  	s.lastWriteRequest = req
    74  	if s.respErr != nil {
    75  		http.Error(w, s.respErr.error, s.respErr.status)
    76  		return
    77  	}
    78  }
    79  
    80  // GetLastWriteRequest returns the last recorded write request.
    81  func (s *TestPromServer) GetLastWriteRequest() *prompb.WriteRequest {
    82  	s.mu.Lock()
    83  	defer s.mu.Unlock()
    84  	return s.lastWriteRequest
    85  }
    86  
    87  // WriteAddr returns http address of a write endpoint.
    88  func (s *TestPromServer) WriteAddr() string {
    89  	return fmt.Sprintf("%s/write", s.svr.URL)
    90  }
    91  
    92  // SetError sets error that will be returned for all incoming requests.
    93  func (s *TestPromServer) SetError(body string, status int) {
    94  	s.mu.Lock()
    95  	defer s.mu.Unlock()
    96  	s.respErr = &respErr{error: body, status: status}
    97  }
    98  
    99  // Reset resets state to default.
   100  func (s *TestPromServer) Reset() {
   101  	s.mu.Lock()
   102  	defer s.mu.Unlock()
   103  	s.respErr = nil
   104  	s.lastWriteRequest = nil
   105  }
   106  
   107  // Close stops underlying http server.
   108  func (s *TestPromServer) Close() {
   109  	s.svr.Close()
   110  }