feat: import assertions when importing store

Author: sujitha-av

cmd/store/create.go

@@ -65,6 +65,11 @@ func CreateStoreWithModel(
 
 	response.Store = *createStoreResponse
 
+	err = fgaClient.SetStoreId(response.Store.Id)
+	if err != nil {
+		return nil, err
+	}
+
 	if inputModel != "" {
 		authModel := authorizationmodel.AuthzModel{}
 

cmd/store/import.go

@@ -123,11 +123,14 @@ func importStore(
 		return nil, err
 	}
 
-	if len(storeData.Tuples) == 0 {
-		return response, nil
+	if len(storeData.Tuples) != 0 {
+		err = importTuples(fgaClient, storeData.Tuples, maxTuplesPerWrite, maxParallelRequests)
+		if err != nil {
+			return nil, err
+		}
 	}
 
-	err = importTuples(fgaClient, storeData.Tuples, maxTuplesPerWrite, maxParallelRequests)
+	err = importAssertions(fgaClient, storeData.Tests, response.Store.Id, response.Model.AuthorizationModelId)
 	if err != nil {
 		return nil, err
 	}
@@ -184,6 +187,47 @@ func importTuples(
 	return nil
 }
 
+func importAssertions(
+	fgaClient client.SdkClient,
+	modelTests []storetest.ModelTest,
+	storeId string,
+	modelId string,
+) error {
+	var assertions []client.ClientAssertion
+
+	for _, modelTest := range modelTests {
+		checkAssertions := getCheckAssertions(modelTest.Check)
+		assertions = append(assertions, checkAssertions...)
+	}
+
+	writeOptions := client.ClientWriteAssertionsOptions{
+		AuthorizationModelId: &modelId,
+		StoreId:              &storeId,
+	}
+
+	if _, err := fgaClient.WriteAssertions(context.Background()).Body(assertions).Options(writeOptions).Execute(); err != nil {
+		return fmt.Errorf("failed to import assertions: %w", err)
+	}
+	return nil
+}
+
+func getCheckAssertions(checkTests []storetest.ModelTestCheck) []client.ClientAssertion {
+	var assertions []client.ClientAssertion
+
+	for _, checkTest := range checkTests {
+		for relation, expectation := range checkTest.Assertions {
+			assertions = append(assertions, client.ClientAssertion{
+				User:        checkTest.User,
+				Relation:    relation,
+				Object:      checkTest.Object,
+				Expectation: expectation,
+			})
+		}
+	}
+
+	return assertions
+}
+
 func createProgressBar(total int) *progressbar.ProgressBar {
 	return progressbar.NewOptions(total,
 		progressbar.OptionSetWriter(os.Stderr),

cmd/store/import_test.go

@@ -0,0 +1,83 @@
+package store
+
+import (
+	"context"
+	"github.com/openfga/cli/internal/fga"
+	mockclient "github.com/openfga/cli/internal/mocks"
+	"github.com/openfga/cli/internal/storetest"
+	"github.com/openfga/go-sdk/client"
+	"go.uber.org/mock/gomock"
+	"testing"
+)
+
+func TestImportStore(t *testing.T) {
+	t.Parallel()
+
+	t.Run("imports assertions into store", func(t *testing.T) {
+		mockCtrl := gomock.NewController(t)
+		defer mockCtrl.Finish()
+
+		mockFgaClient := mockclient.NewMockSdkClient(mockCtrl)
+		clientConfig := fga.ClientConfig{}
+
+		mockWriteAssertions := mockclient.NewMockSdkClientWriteAssertionsRequestInterface(mockCtrl)
+		mockCreateStore := mockclient.NewMockSdkClientCreateStoreRequestInterface(mockCtrl)
+		mockWriteModel := mockclient.NewMockSdkClientWriteAuthorizationModelRequestInterface(mockCtrl)
+
+		expectedAssertions := []client.ClientAssertion{{
+			User:        "user:anne",
+			Relation:    "reader",
+			Object:      "document:doc1",
+			Expectation: true,
+		}}
+
+		modelID := "model-1"
+		storeID := "store-1"
+		expectedOptions := client.ClientWriteAssertionsOptions{
+			AuthorizationModelId: &modelID,
+			StoreId:              &storeID,
+		}
+
+		mockFgaClient.EXPECT().CreateStore(context.Background()).Return(mockCreateStore)
+		mockCreateStore.EXPECT().Body(gomock.Any()).Return(mockCreateStore)
+		mockCreateStore.EXPECT().Execute().Return(&client.ClientCreateStoreResponse{Id: "store-1"}, nil)
+
+		mockFgaClient.EXPECT().SetStoreId("store-1")
+
+		mockFgaClient.EXPECT().WriteAuthorizationModel(context.Background()).Return(mockWriteModel)
+		mockWriteModel.EXPECT().Body(gomock.Any()).Return(mockWriteModel)
+		mockWriteModel.EXPECT().Execute().Return(&client.ClientWriteAuthorizationModelResponse{AuthorizationModelId: "model-1"}, nil)
+
+		mockFgaClient.EXPECT().WriteAssertions(context.Background()).Return(mockWriteAssertions)
+		mockWriteAssertions.EXPECT().Body(expectedAssertions).Return(mockWriteAssertions)
+		mockWriteAssertions.EXPECT().Options(expectedOptions).Return(mockWriteAssertions)
+		mockWriteAssertions.EXPECT().Execute().Return(nil, nil)
+
+		testStore := &storetest.StoreData{
+			Name: "test-store",
+			Model: `type user
+					type document
+						relations
+							define reader: [user]`,
+			Tests: []storetest.ModelTest{
+				{
+					Name: "Test",
+					Check: []storetest.ModelTestCheck{
+						{
+							User:   "user:anne",
+							Object: "document:doc1",
+							Assertions: map[string]bool{
+								"reader": true,
+							},
+						},
+					},
+				},
+			},
+		}
+		_, err := importStore(&clientConfig, mockFgaClient, testStore, "", "", 1, 1, "")
+
+		if err != nil {
+			t.Errorf("expected no error, got %v", err)
+		}
+	})
+}