github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/internal/testserver/server.go (about)

     1  package testserver
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/authzed/spicedb/internal/middleware/servicespecific"
     8  
     9  	"github.com/stretchr/testify/require"
    10  	"google.golang.org/grpc"
    11  
    12  	"github.com/authzed/spicedb/internal/datastore/memdb"
    13  	"github.com/authzed/spicedb/internal/dispatch/graph"
    14  	"github.com/authzed/spicedb/internal/middleware/consistency"
    15  	datastoremw "github.com/authzed/spicedb/internal/middleware/datastore"
    16  	"github.com/authzed/spicedb/pkg/cmd/server"
    17  	"github.com/authzed/spicedb/pkg/cmd/util"
    18  	"github.com/authzed/spicedb/pkg/datastore"
    19  	"github.com/authzed/spicedb/pkg/middleware/logging"
    20  )
    21  
    22  // ServerConfig is configuration for the test server.
    23  type ServerConfig struct {
    24  	MaxUpdatesPerWrite         uint16
    25  	MaxPreconditionsCount      uint16
    26  	MaxRelationshipContextSize int
    27  	StreamingAPITimeout        time.Duration
    28  }
    29  
    30  // NewTestServer creates a new test server, using defaults for the config.
    31  func NewTestServer(require *require.Assertions,
    32  	revisionQuantization time.Duration,
    33  	gcWindow time.Duration,
    34  	schemaPrefixRequired bool,
    35  	dsInitFunc func(datastore.Datastore, *require.Assertions) (datastore.Datastore, datastore.Revision),
    36  ) (*grpc.ClientConn, func(), datastore.Datastore, datastore.Revision) {
    37  	return NewTestServerWithConfig(require, revisionQuantization, gcWindow, schemaPrefixRequired,
    38  		ServerConfig{
    39  			MaxUpdatesPerWrite:         1000,
    40  			MaxPreconditionsCount:      1000,
    41  			StreamingAPITimeout:        30 * time.Second,
    42  			MaxRelationshipContextSize: 25000,
    43  		},
    44  		dsInitFunc)
    45  }
    46  
    47  // NewTestServerWithConfig creates as new test server with the specified config.
    48  func NewTestServerWithConfig(require *require.Assertions,
    49  	revisionQuantization time.Duration,
    50  	gcWindow time.Duration,
    51  	schemaPrefixRequired bool,
    52  	config ServerConfig,
    53  	dsInitFunc func(datastore.Datastore, *require.Assertions) (datastore.Datastore, datastore.Revision),
    54  ) (*grpc.ClientConn, func(), datastore.Datastore, datastore.Revision) {
    55  	emptyDS, err := memdb.NewMemdbDatastore(0, revisionQuantization, gcWindow)
    56  	require.NoError(err)
    57  	ds, revision := dsInitFunc(emptyDS, require)
    58  	ctx, cancel := context.WithCancel(context.Background())
    59  	srv, err := server.NewConfigWithOptions(
    60  		server.WithDatastore(ds),
    61  		server.WithDispatcher(graph.NewLocalOnlyDispatcher(10)),
    62  		server.WithDispatchMaxDepth(50),
    63  		server.WithMaximumPreconditionCount(config.MaxPreconditionsCount),
    64  		server.WithMaximumUpdatesPerWrite(config.MaxUpdatesPerWrite),
    65  		server.WithStreamingAPITimeout(config.StreamingAPITimeout),
    66  		server.WithMaxCaveatContextSize(4096),
    67  		server.WithMaxRelationshipContextSize(config.MaxRelationshipContextSize),
    68  		server.WithGRPCServer(util.GRPCServerConfig{
    69  			Network: util.BufferedNetwork,
    70  			Enabled: true,
    71  		}),
    72  		server.WithSchemaPrefixesRequired(schemaPrefixRequired),
    73  		server.WithGRPCAuthFunc(func(ctx context.Context) (context.Context, error) {
    74  			return ctx, nil
    75  		}),
    76  		server.WithHTTPGateway(util.HTTPServerConfig{HTTPEnabled: false}),
    77  		server.WithMetricsAPI(util.HTTPServerConfig{HTTPEnabled: false}),
    78  		server.WithDispatchServer(util.GRPCServerConfig{Enabled: false}),
    79  		server.SetUnaryMiddlewareModification([]server.MiddlewareModification[grpc.UnaryServerInterceptor]{
    80  			{
    81  				Operation: server.OperationReplaceAllUnsafe,
    82  				Middlewares: []server.ReferenceableMiddleware[grpc.UnaryServerInterceptor]{
    83  					{
    84  						Name:       "logging",
    85  						Middleware: logging.UnaryServerInterceptor(),
    86  					},
    87  					{
    88  						Name:       "datastore",
    89  						Middleware: datastoremw.UnaryServerInterceptor(ds),
    90  					},
    91  					{
    92  						Name:       "consistency",
    93  						Middleware: consistency.UnaryServerInterceptor(),
    94  					},
    95  					{
    96  						Name:       "servicespecific",
    97  						Middleware: servicespecific.UnaryServerInterceptor,
    98  					},
    99  				},
   100  			},
   101  		}),
   102  		server.SetStreamingMiddlewareModification([]server.MiddlewareModification[grpc.StreamServerInterceptor]{
   103  			{
   104  				Operation: server.OperationReplaceAllUnsafe,
   105  				Middlewares: []server.ReferenceableMiddleware[grpc.StreamServerInterceptor]{
   106  					{
   107  						Name:       "logging",
   108  						Middleware: logging.StreamServerInterceptor(),
   109  					},
   110  					{
   111  						Name:       "datastore",
   112  						Middleware: datastoremw.StreamServerInterceptor(ds),
   113  					},
   114  					{
   115  						Name:       "consistency",
   116  						Middleware: consistency.StreamServerInterceptor(),
   117  					},
   118  					{
   119  						Name:       "servicespecific",
   120  						Middleware: servicespecific.StreamServerInterceptor,
   121  					},
   122  				},
   123  			},
   124  		}),
   125  	).Complete(ctx)
   126  	require.NoError(err)
   127  
   128  	go func() {
   129  		require.NoError(srv.Run(ctx))
   130  	}()
   131  
   132  	conn, err := srv.GRPCDialContext(ctx, grpc.WithBlock())
   133  	require.NoError(err)
   134  
   135  	return conn, func() {
   136  		if conn != nil {
   137  			require.NoError(conn.Close())
   138  		}
   139  		cancel()
   140  	}, ds, revision
   141  }