feat: import test tuples and context into assertions

README.md

@@ -274,6 +274,8 @@ fga store **import**
 * `--max-parallel-requests`: Max requests to send in parallel (optional, default=4)
 * `--allow-external-files`: Allow `model_file`, `tuple_file` and `tuple_files` references in the store file to resolve outside the store file's directory (optional, default=false). Only enable this for store files you trust.
 
+The `check` tests in the store file are also written to the store as assertions (at most 100). Each assertion includes the check's `context` and its test's `tuples` (including those loaded from `tuple_file`) as contextual tuples. A test with more than 20 tuples is written without contextual tuples, because assertions accept at most 20, and the CLI prints a warning. If the assertions together are too large for one request, the last ones are written without contextual tuples, and then without context if that's still not enough, also with a warning.
+
 ###### Example
 `fga store import --file model.fga.yaml`
 

cmd/store/import.go

@@ -18,6 +18,7 @@ package store
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"os"
@@ -27,9 +28,12 @@ import (
 
 	"github.com/schollz/progressbar/v3"
 
+	openfgav1 "github.com/openfga/api/proto/openfga/v1"
 	openfga "github.com/openfga/go-sdk"
 	"github.com/openfga/go-sdk/client"
 	"github.com/spf13/cobra"
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/openfga/cli/cmd/model"
 	"github.com/openfga/cli/internal/authorizationmodel"
@@ -46,7 +50,11 @@ const (
 	progressBarThrottleValue = 65
 	progressBarUpdateDelay   = 5 * time.Millisecond
 	maxAssertionsPerWrite    = 100
-	expressionConditionName  = "$expression"
+	// The API accepts at most 20 contextual tuples per assertion.
+	maxContextualTuplesPerAssertion = 20
+	// The API rejects WriteAssertions requests whose assertions add up to more than 64000 bytes in proto encoding.
+	maxAssertionsRequestBytes = 64000
+	expressionConditionName   = "$expression"
 )
 
 var errExpressionConditionNotSupported = errors.New( //nolint:err113
@@ -233,7 +241,17 @@ func importAssertions(
 
 	for _, modelTest := range modelTests {
 		if len(modelTest.Check) > 0 {
-			checkAssertions := getCheckAssertions(modelTest.Check)
+			contextualTuples := modelTest.Tuples
+			if len(contextualTuples) > maxContextualTuplesPerAssertion {
+				fmt.Fprintf(os.Stderr,
+					"Warning: test %q has %d tuples, but assertions accept at most %d contextual tuples, "+
+						"so its assertions are written without them\n",
+					modelTest.Name, len(contextualTuples), maxContextualTuplesPerAssertion)
+
+				contextualTuples = nil
+			}
+
+			checkAssertions := getCheckAssertions(modelTest.Check, contextualTuples)
 			assertions = append(assertions, checkAssertions...)
 		}
 	}
@@ -250,6 +268,8 @@ func importAssertions(
 			assertions = assertions[:maxAssertionsPerWrite]
 		}
 
+		dropContextualTuplesOverSizeLimit(assertions)
+
 		_, err := fgaClient.WriteAssertions(ctx).Body(assertions).Options(writeOptions).Execute()
 		if err != nil {
 			return fmt.Errorf("failed to import test assertions: %w", err)
@@ -259,7 +279,79 @@ func importAssertions(
 	return nil
 }
 
-func getCheckAssertions(checkTests []storetest.ModelTestCheck) []client.ClientAssertion {
+// dropContextualTuplesOverSizeLimit removes the contextual tuples of assertions, starting from the last one,
+// until the assertions fit in one WriteAssertions request. If they still don't fit, it then removes the
+// context of the last assertions too, which main never wrote.
+func dropContextualTuplesOverSizeLimit(assertions []client.ClientAssertion) {
+	sizes := make([]int, len(assertions))
+	total := 0
+
+	for i := range assertions {
+		sizes[i] = assertionSize(assertions[i])
+		total += sizes[i]
+	}
+
+	shrink := func(skip func(client.ClientAssertion) bool, drop func(*client.ClientAssertion)) int {
+		dropped := 0
+
+		for i := len(assertions) - 1; i >= 0 && total > maxAssertionsRequestBytes; i-- {
+			if skip(assertions[i]) {
+				continue
+			}
+
+			drop(&assertions[i])
+			size := assertionSize(assertions[i])
+			total -= sizes[i] - size
+			sizes[i] = size
+			dropped++
+		}
+
+		return dropped
+	}
+
+	droppedTuples := shrink(
+		func(a client.ClientAssertion) bool { return len(a.ContextualTuples) == 0 },
+		func(a *client.ClientAssertion) { a.ContextualTuples = nil },
+	)
+	if droppedTuples > 0 {
+		fmt.Fprintf(os.Stderr,
+			"Warning: the test assertions are over the API's request size limit, "+
+				"so %d of them are written without their test's tuples\n", droppedTuples)
+	}
+
+	droppedContexts := shrink(
+		func(a client.ClientAssertion) bool { return a.Context == nil },
+		func(a *client.ClientAssertion) { a.Context = nil },
+	)
+	if droppedContexts > 0 {
+		fmt.Fprintf(os.Stderr,
+			"Warning: the test assertions are still over the API's request size limit, "+
+				"so %d of them are written without their check's context\n", droppedContexts)
+	}
+}
+
+// assertionSize returns the size of the assertion as the server measures it: the SDK's JSON body decoded
+// into the API's proto message.
+func assertionSize(assertion client.ClientAssertion) int {
+	encoded, err := json.Marshal(assertion.ToAssertion())
+	if err != nil {
+		return 0
+	}
+
+	var protoAssertion openfgav1.Assertion
+	if err := protojson.Unmarshal(encoded, &protoAssertion); err != nil {
+		return len(encoded)
+	}
+
+	return proto.Size(&protoAssertion)
+}
+
+// getCheckAssertions turns check tests into assertions. Each assertion carries the check's context and the
+// test's tuples as contextual tuples, so it holds without writing those tuples to the store.
+func getCheckAssertions(
+	checkTests []storetest.ModelTestCheck,
+	contextualTuples []client.ClientContextualTupleKey,
+) []client.ClientAssertion {
 	assertions := make([]client.ClientAssertion, 0, len(checkTests))
 
 	for _, checkTest := range checkTests {
@@ -270,10 +362,12 @@ func getCheckAssertions(checkTests []storetest.ModelTestCheck) []client.ClientAs
 			for _, object := range objects {
 				for relation, expectation := range checkTest.Assertions {
 					assertions = append(assertions, client.ClientAssertion{
-						User:        user,
-						Relation:    relation,
-						Object:      object,
-						Expectation: expectation,
+						User:             user,
+						Relation:         relation,
+						Object:           object,
+						Expectation:      expectation,
+						Context:          checkTest.Context,
+						ContextualTuples: contextualTuples,
 					})
 				}
 			}

cmd/store/import_test.go

@@ -1,6 +1,7 @@
 package store
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"testing"
@@ -494,3 +495,185 @@ func setupWriteAssertionsMock(
 	mockWriteAssertions.EXPECT().Options(expectedOptions).Return(mockWriteAssertions)
 	mockWriteAssertions.EXPECT().Execute().Return(nil, nil)
 }
+
+func TestImportStoreAssertionsCarryTestTuplesAndContext(t *testing.T) {
+	t.Parallel()
+
+	modelID, storeID := testModelID, testStoreID
+	expectedOptions := client.ClientWriteAssertionsOptions{AuthorizationModelId: &modelID, StoreId: &storeID}
+
+	testTuples := []client.ClientContextualTupleKey{
+		{User: "folder:5", Relation: "parent", Object: "folder:product-2021"},
+	}
+	context := map[string]any{"some": "value"}
+
+	expectedAssertions := []client.ClientAssertion{
+		{
+			User:             "user:anne",
+			Relation:         "viewer",
+			Object:           "folder:product-2021",
+			Expectation:      true,
+			Context:          &context,
+			ContextualTuples: testTuples,
+		},
+	}
+
+	mockCtrl := gomock.NewController(t)
+	mockFgaClient := mockclient.NewMockSdkClient(mockCtrl)
+
+	defer mockCtrl.Finish()
+
+	setupWriteAssertionsMock(mockCtrl, mockFgaClient, expectedAssertions, expectedOptions)
+	setupWriteModelMock(mockCtrl, mockFgaClient, modelID)
+	setupCreateStoreMock(mockCtrl, mockFgaClient, storeID)
+
+	testStore := storetest.StoreData{
+		Model: `type user
+                type folder
+                        relations
+                                define parent: [folder]
+                                define viewer: [user]`,
+		Tests: []storetest.ModelTest{
+			{
+				Name:   "Test",
+				Tuples: testTuples,
+				Check: []storetest.ModelTestCheck{
+					{
+						User:       "user:anne",
+						Object:     "folder:product-2021",
+						Context:    &context,
+						Assertions: map[string]bool{"viewer": true},
+					},
+				},
+			},
+		},
+	}
+
+	_, err := importStore(t.Context(), &fga.ClientConfig{}, mockFgaClient, &testStore, "", "", 10, 1, "")
+	require.NoError(t, err)
+}
+
+func TestImportStoreAssertionsDropTestTuplesOverLimit(t *testing.T) {
+	t.Parallel()
+
+	modelID, storeID := testModelID, testStoreID
+	expectedOptions := client.ClientWriteAssertionsOptions{AuthorizationModelId: &modelID, StoreId: &storeID}
+
+	// One more tuple than an assertion accepts as contextual tuples
+	testTuples := make([]client.ClientContextualTupleKey, maxContextualTuplesPerAssertion+1)
+	for i := range testTuples {
+		testTuples[i] = client.ClientContextualTupleKey{
+			User: fmt.Sprintf("user:%d", i), Relation: "viewer", Object: "document:doc1",
+		}
+	}
+
+	expectedAssertions := []client.ClientAssertion{
+		{User: "user:anne", Relation: "viewer", Object: "document:doc1", Expectation: false},
+	}
+
+	mockCtrl := gomock.NewController(t)
+	mockFgaClient := mockclient.NewMockSdkClient(mockCtrl)
+
+	defer mockCtrl.Finish()
+
+	setupWriteAssertionsMock(mockCtrl, mockFgaClient, expectedAssertions, expectedOptions)
+	setupWriteModelMock(mockCtrl, mockFgaClient, modelID)
+	setupCreateStoreMock(mockCtrl, mockFgaClient, storeID)
+
+	testStore := storetest.StoreData{
+		Model: `type user
+                type document
+                        relations
+                                define viewer: [user]`,
+		Tests: []storetest.ModelTest{
+			{
+				Name:   "Test",
+				Tuples: testTuples,
+				Check: []storetest.ModelTestCheck{
+					{
+						User:       "user:anne",
+						Object:     "document:doc1",
+						Assertions: map[string]bool{"viewer": false},
+					},
+				},
+			},
+		},
+	}
+
+	_, err := importStore(t.Context(), &fga.ClientConfig{}, mockFgaClient, &testStore, "", "", 10, 1, "")
+	require.NoError(t, err)
+}
+
+func TestDropContextualTuplesOverSizeLimit(t *testing.T) {
+	t.Parallel()
+
+	tuples := make([]client.ClientContextualTupleKey, maxContextualTuplesPerAssertion)
+	for i := range tuples {
+		tuples[i] = client.ClientContextualTupleKey{
+			User: fmt.Sprintf("user:member-%02d", i), Relation: "viewer", Object: "document:roadmap-2026",
+		}
+	}
+
+	t.Run("small requests are left alone", func(t *testing.T) {
+		t.Parallel()
+
+		assertions := []client.ClientAssertion{
+			{User: "user:anne", Relation: "viewer", Object: "document:roadmap-2026", ContextualTuples: tuples},
+		}
+		dropContextualTuplesOverSizeLimit(assertions)
+		require.Len(t, assertions[0].ContextualTuples, maxContextualTuplesPerAssertion)
+	})
+
+	t.Run("the last assertions lose their tuples until the request fits", func(t *testing.T) {
+		t.Parallel()
+
+		assertions := make([]client.ClientAssertion, maxAssertionsPerWrite)
+		for i := range assertions {
+			assertions[i] = client.ClientAssertion{
+				User: fmt.Sprintf("user:%d", i), Relation: "viewer", Object: "document:roadmap-2026",
+				ContextualTuples: tuples,
+			}
+		}
+
+		dropContextualTuplesOverSizeLimit(assertions)
+
+		total := 0
+		for _, assertion := range assertions {
+			total += assertionSize(assertion)
+		}
+
+		require.LessOrEqual(t, total, maxAssertionsRequestBytes)
+		require.Len(t, assertions[0].ContextualTuples, maxContextualTuplesPerAssertion)
+		require.Empty(t, assertions[len(assertions)-1].ContextualTuples)
+	})
+
+	t.Run("the last assertions lose their context when dropping tuples isn't enough", func(t *testing.T) {
+		t.Parallel()
+
+		values := make([]any, 60)
+		for i := range values {
+			values[i] = i
+		}
+
+		checkContext := map[string]any{"allowed": values}
+
+		assertions := make([]client.ClientAssertion, maxAssertionsPerWrite)
+		for i := range assertions {
+			assertions[i] = client.ClientAssertion{
+				User: fmt.Sprintf("user:%d", i), Relation: "viewer", Object: "document:roadmap-2026",
+				Context: &checkContext,
+			}
+		}
+
+		dropContextualTuplesOverSizeLimit(assertions)
+
+		total := 0
+		for _, assertion := range assertions {
+			total += assertionSize(assertion)
+		}
+
+		require.LessOrEqual(t, total, maxAssertionsRequestBytes)
+		require.NotNil(t, assertions[0].Context)
+		require.Nil(t, assertions[len(assertions)-1].Context)
+	})
+}