github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/api/api_suite_test.go (about)

     1  package api_test
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"io"
     7  	"net/http"
     8  	"os"
     9  	"testing"
    10  
    11  	"github.com/gorilla/mux"
    12  	. "github.com/onsi/ginkgo/v2"
    13  	. "github.com/onsi/gomega"
    14  
    15  	"github.com/pyroscope-io/pyroscope/pkg/api"
    16  	"github.com/pyroscope-io/pyroscope/pkg/api/router"
    17  	"github.com/pyroscope-io/pyroscope/pkg/model"
    18  	"github.com/pyroscope-io/pyroscope/pkg/server/httputils"
    19  )
    20  
    21  func TestAPI(t *testing.T) {
    22  	RegisterFailHandler(Fail)
    23  	RunSpecs(t, "API Suite")
    24  }
    25  
    26  // requestContextProvider wraps incoming request context.
    27  // Mainly used to inject auth info; use defaultUserCtx if not sure.
    28  type requestContextProvider func(context.Context) context.Context
    29  
    30  var (
    31  	defaultUserCtx = ctxWithUser(&model.User{ID: 1, Role: model.AdminRole})
    32  	defaultReqCtx  = func(ctx context.Context) context.Context { return ctx }
    33  )
    34  
    35  func ctxWithUser(u *model.User) requestContextProvider {
    36  	return func(ctx context.Context) context.Context {
    37  		return model.WithUser(ctx, *u)
    38  	}
    39  }
    40  
    41  func ctxWithAPIKey(k *model.APIKey) requestContextProvider {
    42  	return func(ctx context.Context) context.Context {
    43  		return model.WithAPIKey(ctx, *k)
    44  	}
    45  }
    46  
    47  // newTestRouter initializes http router for testing purposes.
    48  // It was decided to test the whole flow of the request handling,
    49  // including routing, authentication, and authorization.
    50  func newTestRouter(rcp requestContextProvider, services router.Services) *router.Router {
    51  	// For backward compatibility, the redirect handler is invoked
    52  	// if no credentials provided, or the user can not be found.
    53  	// For API key authentication the response code is 401.
    54  	//
    55  	// Note that the handler does not actually redirect but
    56  	// only responds with a distinct code: that's done for
    57  	// testing purposes only.
    58  	redirect := func(w http.ResponseWriter, r *http.Request) {
    59  		services.Logger.WithField("url", r.URL).Debug("redirecting")
    60  		w.WriteHeader(http.StatusTemporaryRedirect)
    61  	}
    62  
    63  	r := router.New(
    64  		mux.NewRouter(),
    65  		services)
    66  
    67  	if services.AuthService != nil {
    68  		r.Use(api.AuthMiddleware(redirect, r.AuthService, httputils.NewDefaultHelper(services.Logger)))
    69  	}
    70  
    71  	r.Use(func(next http.Handler) http.Handler {
    72  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    73  			next.ServeHTTP(w, r.WithContext(rcp(r.Context())))
    74  		})
    75  	})
    76  
    77  	r.RegisterHandlers()
    78  	return r
    79  }
    80  
    81  // newRequest creates an HTTP request with the body specified specified
    82  // as a file name relative to the "testdata".
    83  func newRequest(method, url, body string) *http.Request {
    84  	var reqBody io.Reader
    85  	if body != "" {
    86  		reqBody = readFile(body)
    87  	}
    88  	req, err := http.NewRequest(method, url, reqBody)
    89  	Expect(err).ToNot(HaveOccurred())
    90  	return req
    91  }
    92  
    93  // expectResponse performs an HTTP request and validates the response
    94  // code and body which is specified as a file name relative to the "testdata".
    95  func expectResponse(req *http.Request, body string, code int) {
    96  	response, err := http.DefaultClient.Do(req)
    97  	Expect(err).ToNot(HaveOccurred())
    98  	Expect(response).ToNot(BeNil())
    99  	Expect(response.StatusCode).To(Equal(code))
   100  	if body == "" {
   101  		Expect(readBody(response).String()).To(BeEmpty())
   102  		return
   103  	}
   104  	// It may also make sense to accept the response as a template
   105  	// and render non-deterministic values.
   106  	Expect(readBody(response)).To(MatchJSON(readFile(body)))
   107  }
   108  
   109  func readFile(path string) *bytes.Buffer {
   110  	b, err := os.ReadFile("testdata/" + path)
   111  	Expect(err).ToNot(HaveOccurred())
   112  	return bytes.NewBuffer(b)
   113  }
   114  
   115  func readBody(r *http.Response) *bytes.Buffer {
   116  	b, err := io.ReadAll(r.Body)
   117  	Expect(err).ToNot(HaveOccurred())
   118  	_ = r.Body.Close()
   119  	return bytes.NewBuffer(b)
   120  }