feat: support importing assertions as part of a store (#446)

* refactor: use fgaclient from input param in CreateStoreWithModel

* feat: import assertions when importing store

* refactor and test: import assertions when importing store

* fix linting errors

* test: move clientConfig to be defined per test to avoid false positive race detection

---------

Co-authored-by: Ewan Harris <ewan.harris@okta.com>

Author: sujitha-av

cmd/store/create.go

@@ -27,7 +27,6 @@ import (
 	"github.com/openfga/cli/cmd/model"
 	"github.com/openfga/cli/internal/authorizationmodel"
 	"github.com/openfga/cli/internal/cmdutils"
-	"github.com/openfga/cli/internal/fga"
 	"github.com/openfga/cli/internal/output"
 )
 
@@ -48,16 +47,11 @@ func create(fgaClient client.SdkClient, storeName string) (*client.ClientCreateS
 }
 
 func CreateStoreWithModel(
-	clientConfig fga.ClientConfig,
+	fgaClient client.SdkClient,
 	storeName string,
 	inputModel string,
 	inputFormat authorizationmodel.ModelFormat,
 ) (*CreateStoreAndModelResponse, error) {
-	fgaClient, err := clientConfig.GetFgaClient()
-	if err != nil {
-		return nil, fmt.Errorf("failed to initialize FGA Client due to %w", err)
-	}
-
 	response := CreateStoreAndModelResponse{}
 
 	if storeName == "" {
@@ -73,7 +67,7 @@ func CreateStoreWithModel(
 
 	err = fgaClient.SetStoreId(response.Store.Id)
 	if err != nil {
-		return nil, err //nolint:wrapcheck
+		return nil, fmt.Errorf("failed to set store ID: %w", err)
 	}
 
 	if inputModel != "" {
@@ -109,6 +103,10 @@ export FGA_STORE_ID=$(fga store create --model Model.fga | jq -r .store.id)
 	RunE: func(cmd *cobra.Command, args []string) error {
 		clientConfig := cmdutils.GetClientConfig(cmd)
 		storeName, _ := cmd.Flags().GetString("name")
+		fgaClient, err := clientConfig.GetFgaClient()
+		if err != nil {
+			return fmt.Errorf("failed to initialize FGA Client: %w", err)
+		}
 
 		var inputModel string
 		if err := authorizationmodel.ReadFromInputFileOrArg(
@@ -122,7 +120,7 @@ export FGA_STORE_ID=$(fga store create --model Model.fga | jq -r .store.id)
 			return err //nolint:wrapcheck
 		}
 
-		response, err := CreateStoreWithModel(clientConfig, storeName, inputModel, createModelInputFormat)
+		response, err := CreateStoreWithModel(fgaClient, storeName, inputModel, createModelInputFormat)
 		if err != nil {
 			return err
 		}

cmd/store/import.go

@@ -49,6 +49,7 @@ const (
 // createStore creates a new store with the given client configuration and store data.
 func createStore(
 	clientConfig *fga.ClientConfig,
+	fgaClient client.SdkClient,
 	storeData *storetest.StoreData,
 	format authorizationmodel.ModelFormat,
 	fileName string,
@@ -58,7 +59,7 @@ func createStore(
 		storeDataName = strings.TrimSuffix(path.Base(fileName), ".fga.yaml")
 	}
 
-	createStoreAndModelResponse, err := CreateStoreWithModel(*clientConfig, storeDataName, storeData.Model, format)
+	createStoreAndModelResponse, err := CreateStoreWithModel(fgaClient, storeDataName, storeData.Model, format)
 	if err != nil {
 		return nil, err
 	}
@@ -122,13 +123,18 @@ 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)
-	if err != nil {
-		return nil, err
+	if len(storeData.Tests) != 0 && response.Model != nil {
+		err = importAssertions(fgaClient, storeData.Tests, response.Store.Id, response.Model.AuthorizationModelId)
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	return response, nil
@@ -143,7 +149,7 @@ func createOrUpdateStore(
 	fileName string,
 ) (*CreateStoreAndModelResponse, error) {
 	if storeID == "" {
-		return createStore(clientConfig, storeData, format, fileName)
+		return createStore(clientConfig, fgaClient, storeData, format, fileName)
 	}
 
 	return updateStore(clientConfig, fgaClient, storeData, format, storeID)
@@ -183,6 +189,53 @@ 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 {
+		if len(modelTest.Check) > 0 {
+			checkAssertions := getCheckAssertions(modelTest.Check)
+			assertions = append(assertions, checkAssertions...)
+		}
+	}
+
+	if len(assertions) > 0 {
+		writeOptions := client.ClientWriteAssertionsOptions{
+			AuthorizationModelId: &modelID,
+			StoreId:              &storeID,
+		}
+
+		_, err := fgaClient.WriteAssertions(context.Background()).Body(assertions).Options(writeOptions).Execute()
+		if 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),