From 0aca99d1a4f5a376626700a4e1cfef007d261a75 Mon Sep 17 00:00:00 2001
From: Joseph Schorr <josephschorr@users.noreply.github.com>
Date: Thu, 19 Dec 2024 14:58:30 -0500
Subject: [PATCH] Add caveated bulk load test to datastore tests

---
 internal/testfixtures/generator.go | 21 +++++++++++++++
 pkg/datastore/test/bulk.go         | 43 ++++++++++++++++++++++++++++++
 pkg/datastore/test/datastore.go    |  5 ++--
 3 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/internal/testfixtures/generator.go b/internal/testfixtures/generator.go
index 4e1f578558..9d7417ea84 100644
--- a/internal/testfixtures/generator.go
+++ b/internal/testfixtures/generator.go
@@ -7,7 +7,10 @@ import (
 	"testing"
 	"time"
 
+	"google.golang.org/protobuf/types/known/structpb"
+
 	"github.com/authzed/spicedb/pkg/datastore"
+	corev1 "github.com/authzed/spicedb/pkg/proto/core/v1"
 	"github.com/authzed/spicedb/pkg/tuple"
 )
 
@@ -39,6 +42,7 @@ func NewBulkRelationshipGenerator(objectType, relation, subjectType string, coun
 		relation,
 		subjectType,
 		false,
+		false,
 	}
 }
 
@@ -49,6 +53,7 @@ type BulkRelationshipGenerator struct {
 	relation       string
 	subjectType    string
 	WithExpiration bool
+	WithCaveat     bool
 }
 
 func (btg *BulkRelationshipGenerator) Next(_ context.Context) (*tuple.Relationship, error) {
@@ -63,6 +68,21 @@ func (btg *BulkRelationshipGenerator) Next(_ context.Context) (*tuple.Relationsh
 		expiration = &exp
 	}
 
+	var caveat *corev1.ContextualizedCaveat
+	if btg.WithCaveat {
+		c, err := structpb.NewStruct(map[string]interface{}{
+			"secret": "1235",
+		})
+		if err != nil {
+			btg.t.Fatalf("failed to create struct: %v", err)
+		}
+
+		caveat = &corev1.ContextualizedCaveat{
+			CaveatName: "test",
+			Context:    c,
+		}
+	}
+
 	return &tuple.Relationship{
 		RelationshipReference: tuple.RelationshipReference{
 			Resource: tuple.ObjectAndRelation{
@@ -76,6 +96,7 @@ func (btg *BulkRelationshipGenerator) Next(_ context.Context) (*tuple.Relationsh
 				Relation:   datastore.Ellipsis,
 			},
 		},
+		OptionalCaveat:     caveat,
 		OptionalExpiration: expiration,
 	}, nil
 }
diff --git a/pkg/datastore/test/bulk.go b/pkg/datastore/test/bulk.go
index c9500e5cf9..ca0f64a3b7 100644
--- a/pkg/datastore/test/bulk.go
+++ b/pkg/datastore/test/bulk.go
@@ -122,6 +122,49 @@ func BulkUploadAlreadyExistsSameCallErrorTest(t *testing.T, tester DatastoreTest
 	grpcutil.RequireStatus(t, codes.AlreadyExists, err)
 }
 
+func BulkUploadWithCaveats(t *testing.T, tester DatastoreTester) {
+	tc := 10
+	require := require.New(t)
+	ctx := context.Background()
+
+	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
+	require.NoError(err)
+
+	ds, _ := testfixtures.StandardDatastoreWithSchema(rawDS, require)
+	bulkSource := testfixtures.NewBulkRelationshipGenerator(
+		testfixtures.DocumentNS.Name,
+		"caveated_viewer",
+		testfixtures.UserNS.Name,
+		tc,
+		t,
+	)
+	bulkSource.WithCaveat = true
+
+	// This is statically defined so we can cast straight.
+	uintTc, _ := safecast.ToUint64(tc)
+	lastRevision, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
+		loaded, err := rwt.BulkLoad(ctx, bulkSource)
+		require.NoError(err)
+		require.Equal(uintTc, loaded)
+		return err
+	})
+	require.NoError(err)
+
+	iter, err := ds.SnapshotReader(lastRevision).QueryRelationships(ctx, datastore.RelationshipsFilter{
+		OptionalResourceType: testfixtures.DocumentNS.Name,
+	})
+	require.NoError(err)
+
+	for found, err := range iter {
+		require.NoError(err)
+		require.Nil(found.OptionalExpiration)
+		require.NotNil(found.OptionalCaveat)
+		require.NotEmpty(found.OptionalCaveat.CaveatName)
+		require.NotNil(found.OptionalCaveat.Context)
+		require.Equal(map[string]any{"secret": "1235"}, found.OptionalCaveat.Context.AsMap())
+	}
+}
+
 func BulkUploadWithExpiration(t *testing.T, tester DatastoreTester) {
 	tc := 10
 	require := require.New(t)
diff --git a/pkg/datastore/test/datastore.go b/pkg/datastore/test/datastore.go
index 4e58eb4562..54a898db26 100644
--- a/pkg/datastore/test/datastore.go
+++ b/pkg/datastore/test/datastore.go
@@ -161,8 +161,9 @@ func AllWithExceptions(t *testing.T, tester DatastoreTester, except Categories,
 	t.Run("TestBulkUploadErrors", runner(tester, BulkUploadErrorsTest))
 	t.Run("TestBulkUploadAlreadyExistsError", runner(tester, BulkUploadAlreadyExistsErrorTest))
 	t.Run("TestBulkUploadAlreadyExistsSameCallError", runner(tester, BulkUploadAlreadyExistsSameCallErrorTest))
-	t.Run("BulkUploadEditCaveat", runner(tester, BulkUploadEditCaveat))
-	t.Run("BulkUploadWithExpiration", runner(tester, BulkUploadWithExpiration))
+	t.Run("TestBulkUploadEditCaveat", runner(tester, BulkUploadEditCaveat))
+	t.Run("TestBulkUploadWithCaveats", runner(tester, BulkUploadWithCaveats))
+	t.Run("TestBulkUploadWithExpiration", runner(tester, BulkUploadWithExpiration))
 
 	if !except.Stats() {
 		t.Run("TestStats", runner(tester, StatsTest))