github.com/sequix/cortex@v1.1.6/pkg/chunk/gcp/fixtures.go (about)

     1  package gcp
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"cloud.google.com/go/bigtable"
     8  	"cloud.google.com/go/bigtable/bttest"
     9  	"github.com/fsouza/fake-gcs-server/fakestorage"
    10  	"google.golang.org/api/option"
    11  	"google.golang.org/grpc"
    12  
    13  	"github.com/sequix/cortex/pkg/chunk"
    14  	"github.com/sequix/cortex/pkg/chunk/testutils"
    15  )
    16  
    17  const (
    18  	proj, instance = "proj", "instance"
    19  )
    20  
    21  type fixture struct {
    22  	btsrv  *bttest.Server
    23  	gcssrv *fakestorage.Server
    24  
    25  	name string
    26  
    27  	gcsObjectClient bool
    28  	columnKeyClient bool
    29  	hashPrefix      bool
    30  }
    31  
    32  func (f *fixture) Name() string {
    33  	return f.name
    34  }
    35  
    36  func (f *fixture) Clients() (
    37  	iClient chunk.IndexClient, cClient chunk.ObjectClient, tClient chunk.TableClient,
    38  	schemaConfig chunk.SchemaConfig, err error,
    39  ) {
    40  	f.btsrv, err = bttest.NewServer("localhost:0")
    41  	if err != nil {
    42  		return
    43  	}
    44  
    45  	f.gcssrv = fakestorage.NewServer(nil)
    46  	f.gcssrv.CreateBucket("chunks")
    47  
    48  	conn, err := grpc.Dial(f.btsrv.Addr, grpc.WithInsecure())
    49  	if err != nil {
    50  		return
    51  	}
    52  
    53  	ctx := context.Background()
    54  	adminClient, err := bigtable.NewAdminClient(ctx, proj, instance, option.WithGRPCConn(conn))
    55  	if err != nil {
    56  		return
    57  	}
    58  
    59  	schemaConfig = testutils.DefaultSchemaConfig("gcp-columnkey")
    60  	tClient = &tableClient{
    61  		client: adminClient,
    62  	}
    63  
    64  	client, err := bigtable.NewClient(ctx, proj, instance, option.WithGRPCConn(conn))
    65  	if err != nil {
    66  		return
    67  	}
    68  
    69  	cfg := Config{
    70  		DistributeKeys: f.hashPrefix,
    71  	}
    72  	if f.columnKeyClient {
    73  		iClient = newStorageClientColumnKey(cfg, schemaConfig, client)
    74  	} else {
    75  		iClient = newStorageClientV1(cfg, schemaConfig, client)
    76  	}
    77  
    78  	if f.gcsObjectClient {
    79  		cClient = newGCSObjectClient(GCSConfig{
    80  			BucketName: "chunks",
    81  		}, schemaConfig, f.gcssrv.Client())
    82  	} else {
    83  		cClient = newBigtableObjectClient(Config{}, schemaConfig, client)
    84  	}
    85  
    86  	return
    87  }
    88  
    89  func (f *fixture) Teardown() error {
    90  	f.btsrv.Close()
    91  	f.gcssrv.Stop()
    92  	return nil
    93  }
    94  
    95  // Fixtures for unit testing GCP storage.
    96  var Fixtures = func() []testutils.Fixture {
    97  	fixtures := []testutils.Fixture{}
    98  	for _, gcsObjectClient := range []bool{true, false} {
    99  		for _, columnKeyClient := range []bool{true, false} {
   100  			for _, hashPrefix := range []bool{true, false} {
   101  				fixtures = append(fixtures, &fixture{
   102  					name:            fmt.Sprintf("bigtable-columnkey:%v-gcsObjectClient:%v-hashPrefix:%v", columnKeyClient, gcsObjectClient, hashPrefix),
   103  					columnKeyClient: columnKeyClient,
   104  					gcsObjectClient: gcsObjectClient,
   105  					hashPrefix:      hashPrefix,
   106  				})
   107  			}
   108  		}
   109  	}
   110  	return fixtures
   111  }()