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

     1  package testfixtures
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"strconv"
     7  	"testing"
     8  
     9  	"github.com/authzed/spicedb/pkg/datastore"
    10  	core "github.com/authzed/spicedb/pkg/proto/core/v1"
    11  )
    12  
    13  const (
    14  	FirstLetters      = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890_"
    15  	SubsequentLetters = FirstLetters + "/_|-"
    16  )
    17  
    18  func RandomObjectID(length uint8) string {
    19  	b := make([]byte, length)
    20  	for i := range b {
    21  		sourceLetters := SubsequentLetters
    22  		if i == 0 {
    23  			sourceLetters = FirstLetters
    24  		}
    25  		// nolint:gosec
    26  		// G404 use of non cryptographically secure random number generator is not a security concern here,
    27  		// as this is only used for generating fixtures in testing.
    28  		b[i] = sourceLetters[rand.Intn(len(sourceLetters))]
    29  	}
    30  	return string(b)
    31  }
    32  
    33  func NewBulkTupleGenerator(objectType, relation, subjectType string, count int, t *testing.T) *BulkTupleGenerator {
    34  	return &BulkTupleGenerator{
    35  		count,
    36  		t,
    37  		core.RelationTuple{
    38  			ResourceAndRelation: &core.ObjectAndRelation{
    39  				Namespace: objectType,
    40  				Relation:  relation,
    41  			},
    42  			Subject: &core.ObjectAndRelation{
    43  				Namespace: subjectType,
    44  				Relation:  datastore.Ellipsis,
    45  			},
    46  		},
    47  	}
    48  }
    49  
    50  type BulkTupleGenerator struct {
    51  	remaining int
    52  	t         *testing.T
    53  
    54  	current core.RelationTuple
    55  }
    56  
    57  func (btg *BulkTupleGenerator) Next(_ context.Context) (*core.RelationTuple, error) {
    58  	if btg.remaining <= 0 {
    59  		return nil, nil
    60  	}
    61  	btg.remaining--
    62  	btg.current.ResourceAndRelation.ObjectId = strconv.Itoa(btg.remaining)
    63  	btg.current.Subject.ObjectId = strconv.Itoa(btg.remaining)
    64  
    65  	return &btg.current, nil
    66  }
    67  
    68  var _ datastore.BulkWriteRelationshipSource = &BulkTupleGenerator{}