go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/testing/prpctest/server.go (about)

     1  // Copyright 2016 The LUCI Authors.
     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 prpctest is a package to facilitate pRPC testing by wrapping
    16  // httptest with a pRPC Server.
    17  package prpctest
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"net/http/httptest"
    23  	"net/url"
    24  
    25  	"go.chromium.org/luci/grpc/prpc"
    26  	"go.chromium.org/luci/server/router"
    27  )
    28  
    29  // Server is a pRPC test server.
    30  type Server struct {
    31  	prpc.Server
    32  
    33  	// Base returns a middleware chain. It is handed the Context passed to
    34  	// Start. If Base is nil, setContext will be used.
    35  	Base func(context.Context) router.MiddlewareChain
    36  
    37  	// HTTP is the active HTTP test server. It will be valid when the Server is
    38  	// running.
    39  	HTTP *httptest.Server
    40  
    41  	// Host is the server address ("addr:port") if it is running.
    42  	Host string
    43  }
    44  
    45  func setContext(ctx context.Context) router.MiddlewareChain {
    46  	return router.NewMiddlewareChain(
    47  		func(rctx *router.Context, next router.Handler) {
    48  			rctx.Request = rctx.Request.WithContext(ctx)
    49  			next(rctx)
    50  		},
    51  	)
    52  }
    53  
    54  // Start starts the server. Any currently-registered services will be installed
    55  // into the pRPC Server.
    56  func (s *Server) Start(ctx context.Context) {
    57  	// Clean up any active server.
    58  	s.Close()
    59  
    60  	base := s.Base
    61  	if base == nil {
    62  		base = setContext
    63  	}
    64  
    65  	r := router.New()
    66  	s.InstallHandlers(r, base(ctx))
    67  	s.HTTP = httptest.NewServer(r)
    68  
    69  	u, err := url.Parse(s.HTTP.URL)
    70  	if err != nil {
    71  		panic(err)
    72  	}
    73  	s.Host = u.Host
    74  }
    75  
    76  // NewClient returns a prpc.Client configured to use the Server.
    77  func (s *Server) NewClient() (*prpc.Client, error) {
    78  	return s.NewClientWithOptions(nil)
    79  }
    80  
    81  // NewClientWithOptions returns a prpc.Client configured to use the Server.
    82  func (s *Server) NewClientWithOptions(opts *prpc.Options) (*prpc.Client, error) {
    83  	if s.HTTP == nil {
    84  		return nil, errors.New("not running")
    85  	}
    86  
    87  	if opts == nil {
    88  		opts = &prpc.Options{Insecure: true}
    89  	} else {
    90  		cpy := *opts
    91  		opts = &cpy
    92  		opts.Insecure = true
    93  	}
    94  
    95  	return &prpc.Client{
    96  		Host:    s.Host,
    97  		Options: opts,
    98  	}, nil
    99  }
   100  
   101  // Close closes the Server, releasing any retained resources.
   102  func (s *Server) Close() {
   103  	if s.HTTP != nil {
   104  		s.HTTP.Close()
   105  
   106  		s.HTTP = nil
   107  	}
   108  }