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

     1  package testutils
     2  
     3  import (
     4  	"context"
     5  	"strconv"
     6  	"time"
     7  
     8  	promchunk "github.com/sequix/cortex/pkg/chunk/encoding"
     9  	"github.com/sequix/cortex/pkg/util/flagext"
    10  	"github.com/prometheus/common/model"
    11  	"github.com/prometheus/prometheus/pkg/labels"
    12  
    13  	"github.com/sequix/cortex/pkg/chunk"
    14  	"github.com/sequix/cortex/pkg/ingester/client"
    15  )
    16  
    17  const (
    18  	userID = "userID"
    19  )
    20  
    21  // Fixture type for per-backend testing.
    22  type Fixture interface {
    23  	Name() string
    24  	Clients() (chunk.IndexClient, chunk.ObjectClient, chunk.TableClient, chunk.SchemaConfig, error)
    25  	Teardown() error
    26  }
    27  
    28  // DefaultSchemaConfig returns default schema for use in test fixtures
    29  func DefaultSchemaConfig(kind string) chunk.SchemaConfig {
    30  	schemaConfig := chunk.DefaultSchemaConfig(kind, "v1", model.Now().Add(-time.Hour*2))
    31  	return schemaConfig
    32  }
    33  
    34  // Setup a fixture with initial tables
    35  func Setup(fixture Fixture, tableName string) (chunk.IndexClient, chunk.ObjectClient, error) {
    36  	var tbmConfig chunk.TableManagerConfig
    37  	flagext.DefaultValues(&tbmConfig)
    38  	indexClient, objectClient, tableClient, schemaConfig, err := fixture.Clients()
    39  	if err != nil {
    40  		return nil, nil, err
    41  	}
    42  
    43  	tableManager, err := chunk.NewTableManager(tbmConfig, schemaConfig, 12*time.Hour, tableClient, nil)
    44  	if err != nil {
    45  		return nil, nil, err
    46  	}
    47  
    48  	err = tableManager.SyncTables(context.Background())
    49  	if err != nil {
    50  		return nil, nil, err
    51  	}
    52  
    53  	err = tableClient.CreateTable(context.Background(), chunk.TableDesc{
    54  		Name: tableName,
    55  	})
    56  	return indexClient, objectClient, err
    57  }
    58  
    59  // CreateChunks creates some chunks for testing
    60  func CreateChunks(startIndex, batchSize int, start model.Time) ([]string, []chunk.Chunk, error) {
    61  	keys := []string{}
    62  	chunks := []chunk.Chunk{}
    63  	for j := 0; j < batchSize; j++ {
    64  		chunk := dummyChunkFor(start, labels.Labels{
    65  			{Name: model.MetricNameLabel, Value: "foo"},
    66  			{Name: "index", Value: strconv.Itoa(startIndex*batchSize + j)},
    67  		})
    68  		chunks = append(chunks, chunk)
    69  		keys = append(keys, chunk.ExternalKey())
    70  	}
    71  	return keys, chunks, nil
    72  }
    73  
    74  func dummyChunk(now model.Time) chunk.Chunk {
    75  	return dummyChunkFor(now, labels.Labels{
    76  		{Name: model.MetricNameLabel, Value: "foo"},
    77  		{Name: "bar", Value: "baz"},
    78  		{Name: "toms", Value: "code"},
    79  	})
    80  }
    81  
    82  func dummyChunkFor(now model.Time, metric labels.Labels) chunk.Chunk {
    83  	cs, _ := promchunk.New().Add(model.SamplePair{Timestamp: now, Value: 0})
    84  	chunk := chunk.NewChunk(
    85  		userID,
    86  		client.Fingerprint(metric),
    87  		metric,
    88  		cs[0],
    89  		now.Add(-time.Hour),
    90  		now,
    91  	)
    92  	// Force checksum calculation.
    93  	err := chunk.Encode()
    94  	if err != nil {
    95  		panic(err)
    96  	}
    97  	return chunk
    98  }