Add some additional datastore tests to improve coverage

pkg/datastore/test/counters.go

@@ -2,6 +2,8 @@ package test
 
 import (
 	"context"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -65,6 +67,41 @@ func RelationshipCounterOverExpiredTest(t *testing.T, tester DatastoreTester) {
 	require.Equal(t, expectedCount, count)
 }
 
+func RegisterRelationshipCountersInParallelTest(t *testing.T, tester DatastoreTester) {
+	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
+	require.NoError(t, err)
+
+	ds, _ := testfixtures.StandardDatastoreWithData(rawDS, require.New(t))
+
+	// Run multiple registrations of the counter in parallel and ensure only
+	// one succeeds.
+	var numSucceeded atomic.Int32
+	var numFailed atomic.Int32
+	var wg sync.WaitGroup
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+		go func() {
+			_, err := ds.ReadWriteTx(context.Background(), func(ctx context.Context, tx datastore.ReadWriteTransaction) error {
+				return tx.RegisterCounter(ctx, "document", &core.RelationshipFilter{
+					ResourceType: testfixtures.DocumentNS.Name,
+				})
+			})
+			if err != nil {
+				require.Contains(t, err.Error(), "counter with name `document` already registered")
+				numFailed.Add(1)
+			} else {
+				numSucceeded.Add(1)
+			}
+			wg.Done()
+		}()
+	}
+
+	// Wait for all goroutines to finish.
+	wg.Wait()
+	require.Equal(t, int32(1), numSucceeded.Load())
+	require.Equal(t, int32(9), numFailed.Load())
+}
+
 func RelationshipCountersTest(t *testing.T, tester DatastoreTester) {
 	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
 	require.NoError(t, err)

pkg/datastore/test/datastore.go

@@ -130,11 +130,13 @@ func AllWithExceptions(t *testing.T, tester DatastoreTester, except Categories,
 	t.Run("TestBulkDeleteRelationships", runner(tester, BulkDeleteRelationshipsTest))
 	t.Run("TestDeleteCaveatedTuple", runner(tester, DeleteCaveatedTupleTest))
 	t.Run("TestDeleteWithLimit", runner(tester, DeleteWithLimitTest))
+	t.Run("TestDeleteWithInvalidPrefix", runner(tester, DeleteWithInvalidPrefixTest))
 	t.Run("TestQueryRelationshipsWithVariousFilters", runner(tester, QueryRelationshipsWithVariousFiltersTest))
 	t.Run("TestDeleteRelationshipsWithVariousFilters", runner(tester, DeleteRelationshipsWithVariousFiltersTest))
 	t.Run("TestTouchTypedAlreadyExistingWithoutCaveat", runner(tester, TypedTouchAlreadyExistingTest))
 	t.Run("TestTouchTypedAlreadyExistingWithCaveat", runner(tester, TypedTouchAlreadyExistingWithCaveatTest))
 	t.Run("TestRelationshipExpiration", runner(tester, RelationshipExpirationTest))
+	t.Run("TestMixedWriteOperations", runner(tester, MixedWriteOperationsTest))
 
 	t.Run("TestMultipleReadsInRWT", runner(tester, MultipleReadsInRWTTest))
 	t.Run("TestConcurrentWriteSerialization", runner(tester, ConcurrentWriteSerializationTest))
@@ -198,6 +200,7 @@ func AllWithExceptions(t *testing.T, tester DatastoreTester, except Categories,
 	t.Run("TestUpdateRelationshipCounter", runner(tester, UpdateRelationshipCounterTest))
 	t.Run("TestDeleteAllData", runner(tester, DeleteAllDataTest))
 	t.Run("TestRelationshipCounterOverExpired", runner(tester, RelationshipCounterOverExpiredTest))
+	t.Run("TestRegisterRelationshipCountersInParallel", runner(tester, RegisterRelationshipCountersInParallelTest))
 }
 
 func OnlyGCTests(t *testing.T, tester DatastoreTester, concurrent bool) {

pkg/datastore/test/relationships.go

@@ -687,6 +687,116 @@ func DeleteOneThousandIndividualInOneCallTest(t *testing.T, tester DatastoreTest
 	ensureRelationships(ctx, require, ds, makeTestRel("foo", "extra"))
 }
 
+// DeleteWithInvalidPrefixTest tests deleting relationships with an invalid object prefix.
+func DeleteWithInvalidPrefixTest(t *testing.T, tester DatastoreTester) {
+	require := require.New(t)
+
+	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
+	require.NoError(err)
+
+	ds, _ := testfixtures.StandardDatastoreWithSchema(rawDS, require)
+	ctx := context.Background()
+
+	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
+		_, err := rwt.DeleteRelationships(ctx, &v1.RelationshipFilter{
+			OptionalResourceIdPrefix: "hithere%",
+		})
+		return err
+	})
+	require.Error(err)
+	require.ErrorContains(err, "value does not match regex")
+}
+
+// MixedWriteOperationsTest tests a WriteRelationships call with mixed operations.
+func MixedWriteOperationsTest(t *testing.T, tester DatastoreTester) {
+	require := require.New(t)
+
+	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
+	require.NoError(err)
+
+	ds, _ := testfixtures.StandardDatastoreWithSchema(rawDS, require)
+	ctx := context.Background()
+
+	// Write the 100 relationships.
+	rels := make([]tuple.Relationship, 0, 100)
+	for i := 0; i < 100; i++ {
+		rels = append(rels, tuple.Relationship{
+			RelationshipReference: tuple.RelationshipReference{
+				Resource: tuple.ONR("document", "somedoc", "viewer"),
+				Subject:  tuple.ONR("user", "user-"+strconv.Itoa(i), tuple.Ellipsis),
+			},
+		})
+	}
+
+	_, err = common.WriteRelationships(ctx, ds, tuple.UpdateOperationCreate, rels...)
+	require.NoError(err)
+	ensureRelationships(ctx, require, ds, rels...)
+
+	// Create a WriteRelationships call with a few CREATEs, TOUCHes and DELETEs.
+	updates := make([]tuple.RelationshipUpdate, 0, 30)
+	expectedRels := make([]tuple.Relationship, 0, 105)
+
+	// Add a CREATE for 10 new relationships.
+	for i := 0; i < 10; i++ {
+		newRel := tuple.Relationship{
+			RelationshipReference: tuple.RelationshipReference{
+				Resource: tuple.ONR("document", "somedoc", "viewer"),
+				Subject:  tuple.ONR("user", "user-"+strconv.Itoa(i+100), tuple.Ellipsis),
+			},
+		}
+		expectedRels = append(expectedRels, newRel)
+		updates = append(updates, tuple.Create(newRel))
+	}
+
+	// Add a TOUCH for 5 existing relationships.
+	for i := 0; i < 5; i++ {
+		updates = append(updates, tuple.Touch(tuple.Relationship{
+			RelationshipReference: tuple.RelationshipReference{
+				Resource: tuple.ONR("document", "somedoc", "viewer"),
+				Subject:  tuple.ONR("user", "user-"+strconv.Itoa(i+50), tuple.Ellipsis),
+			},
+		}))
+	}
+
+	// Add a TOUCH for 5 new relationships.
+	for i := 0; i < 5; i++ {
+		newRel := tuple.Relationship{
+			RelationshipReference: tuple.RelationshipReference{
+				Resource: tuple.ONR("document", "somedoc", "viewer"),
+				Subject:  tuple.ONR("user", "user-"+strconv.Itoa(i+110), tuple.Ellipsis),
+			},
+		}
+		expectedRels = append(expectedRels, newRel)
+		updates = append(updates, tuple.Touch(newRel))
+	}
+
+	// DELETE the first 10 relationships.
+	deletedRels := make([]tuple.Relationship, 0, 10)
+	for i := 0; i < 10; i++ {
+		rel := tuple.Relationship{
+			RelationshipReference: tuple.RelationshipReference{
+				Resource: tuple.ONR("document", "somedoc", "viewer"),
+				Subject:  tuple.ONR("user", "user-"+strconv.Itoa(i), tuple.Ellipsis),
+			},
+		}
+		deletedRels = append(deletedRels, rel)
+		updates = append(updates, tuple.Delete(rel))
+	}
+
+	for i := 10; i < 100; i++ {
+		expectedRels = append(expectedRels, rels[i])
+	}
+
+	_, err = ds.ReadWriteTx(ctx, func(ctx context.Context, rwt datastore.ReadWriteTransaction) error {
+		return rwt.WriteRelationships(ctx, updates)
+	})
+	require.NoError(err)
+
+	// Ensure the expected relationships all exist.
+	ensureRelationships(ctx, require, ds, expectedRels...)
+	ensureNotRelationships(ctx, require, ds, deletedRels...)
+}
+
 // DeleteWithLimitTest tests deleting relationships with a limit.
 func DeleteWithLimitTest(t *testing.T, tester DatastoreTester) {
 	require := require.New(t)