chore: Clean up old test object helpers (#3049)

- remove old helpers (testDb, testSchema, testWarehouse) - use `TestClient` instead of `Client` in the acceptance tests - add missing helpers - add missing helper methods - move grant helpers from role client to grant client
Snowflake-Labs · Sep 10, 2024 · 78ecb62 · 78ecb62
1 parent 5f54ff9
commit 78ecb62
Show file tree

Hide file tree

Showing 46 changed files with 636 additions and 684 deletions.
diff --git a/pkg/acceptance/bettertestspoc/assert/objectassert/view_snowflake_ext.go b/pkg/acceptance/bettertestspoc/assert/objectassert/view_snowflake_ext.go
@@ -51,12 +51,12 @@ func (v *ViewAssert) HasNoProjectionPolicyReferences(client *helpers.TestClient)
 func (v *ViewAssert) hasNoPolicyReference(client *helpers.TestClient, kind sdk.PolicyKind) *ViewAssert {
 	v.AddAssertion(func(t *testing.T, o *sdk.View) error {
 		t.Helper()
-		refs, err := client.PolicyReferences.GetPolicyReferences(t, o.ID(), sdk.ObjectTypeView)
+		refs, err := client.PolicyReferences.GetPolicyReferences(t, o.ID(), sdk.PolicyEntityDomainView)
 		if err != nil {
 			return err
 		}
-		refs = slices.DeleteFunc(refs, func(reference helpers.PolicyReference) bool {
-			return reference.PolicyKind != string(kind)
+		refs = slices.DeleteFunc(refs, func(reference sdk.PolicyReference) bool {
+			return reference.PolicyKind != kind
 		})
 		if len(refs) > 0 {
 			return fmt.Errorf("expected no %s policy references; got: %v", kind, refs)
@@ -85,12 +85,12 @@ func (v *ViewAssert) HasProjectionPolicyReferences(client *helpers.TestClient, n
 func (v *ViewAssert) hasPolicyReference(client *helpers.TestClient, kind sdk.PolicyKind, n int) *ViewAssert {
 	v.AddAssertion(func(t *testing.T, o *sdk.View) error {
 		t.Helper()
-		refs, err := client.PolicyReferences.GetPolicyReferences(t, o.ID(), sdk.ObjectTypeView)
+		refs, err := client.PolicyReferences.GetPolicyReferences(t, o.ID(), sdk.PolicyEntityDomainView)
 		if err != nil {
 			return err
 		}
-		refs = slices.DeleteFunc(refs, func(reference helpers.PolicyReference) bool {
-			return reference.PolicyKind != string(kind)
+		refs = slices.DeleteFunc(refs, func(reference sdk.PolicyReference) bool {
+			return reference.PolicyKind != kind
 		})
 		if len(refs) != n {
 			return fmt.Errorf("expected %d %s policy references; got: %d, %v", n, kind, len(refs), refs)

diff --git a/pkg/acceptance/check_destroy.go b/pkg/acceptance/check_destroy.go
@@ -17,7 +17,8 @@ import (
 
 func CheckDestroy(t *testing.T, resource resources.Resource) func(*terraform.State) error {
 	t.Helper()
-	client := Client(t)
+	// TODO [SNOW-1653619]: use TestClient() here
+	client := atc.client
 	t.Logf("running check destroy for resource %s", resource)
 
 	return func(s *terraform.State) error {
@@ -236,24 +237,18 @@ func asId[T sdk.AccountObjectIdentifier | sdk.DatabaseObjectIdentifier | sdk.Sch
 // CheckGrantAccountRoleDestroy is a custom checks that should be later incorporated into generic CheckDestroy
 func CheckGrantAccountRoleDestroy(t *testing.T) func(*terraform.State) error {
 	t.Helper()
-	client := Client(t)
 
 	return func(s *terraform.State) error {
 		for _, rs := range s.RootModule().Resources {
 			if rs.Type != "snowflake_grant_account_role" {
 				continue
 			}
-			ctx := context.Background()
 			parts := strings.Split(rs.Primary.ID, "|")
 			roleName := parts[0]
 			roleIdentifier := sdk.NewAccountObjectIdentifierFromFullyQualifiedName(roleName)
 			objectType := parts[1]
 			targetIdentifier := parts[2]
-			grants, err := client.Grants.Show(ctx, &sdk.ShowGrantOptions{
-				Of: &sdk.ShowGrantsOf{
-					Role: roleIdentifier,
-				},
-			})
+			grants, err := TestClient().Grant.ShowGrantsOfAccountRole(t, roleIdentifier)
 			if err != nil {
 				return nil
 			}
@@ -278,24 +273,18 @@ func CheckGrantAccountRoleDestroy(t *testing.T) func(*terraform.State) error {
 // CheckGrantDatabaseRoleDestroy is a custom checks that should be later incorporated into generic CheckDestroy
 func CheckGrantDatabaseRoleDestroy(t *testing.T) func(*terraform.State) error {
 	t.Helper()
-	client := Client(t)
 
 	return func(s *terraform.State) error {
 		for _, rs := range s.RootModule().Resources {
 			if rs.Type != "snowflake_grant_database_role" {
 				continue
 			}
-			ctx := context.Background()
 			id := rs.Primary.ID
 			ids := strings.Split(id, "|")
 			databaseRoleName := ids[0]
 			objectType := ids[1]
 			parentRoleName := ids[2]
-			grants, err := client.Grants.Show(ctx, &sdk.ShowGrantOptions{
-				Of: &sdk.ShowGrantsOf{
-					DatabaseRole: sdk.NewDatabaseObjectIdentifierFromFullyQualifiedName(databaseRoleName),
-				},
-			})
+			grants, err := TestClient().Grant.ShowGrantsOfDatabaseRole(t, sdk.NewDatabaseObjectIdentifierFromFullyQualifiedName(databaseRoleName))
 			if err != nil {
 				continue
 			}
@@ -314,21 +303,15 @@ func CheckGrantDatabaseRoleDestroy(t *testing.T) func(*terraform.State) error {
 // CheckAccountRolePrivilegesRevoked is a custom checks that should be later incorporated into generic CheckDestroy
 func CheckAccountRolePrivilegesRevoked(t *testing.T) func(*terraform.State) error {
 	t.Helper()
-	client := Client(t)
 
 	return func(s *terraform.State) error {
 		for _, rs := range s.RootModule().Resources {
 			if rs.Type != "snowflake_grant_privileges_to_account_role" {
 				continue
 			}
-			ctx := context.Background()
 
 			id := sdk.NewAccountObjectIdentifierFromFullyQualifiedName(rs.Primary.Attributes["account_role_name"])
-			grants, err := client.Grants.Show(ctx, &sdk.ShowGrantOptions{
-				To: &sdk.ShowGrantsTo{
-					Role: id,
-				},
-			})
+			grants, err := TestClient().Grant.ShowGrantsToAccountRole(t, id)
 			if err != nil {
 				if errors.Is(err, sdk.ErrObjectNotExistOrAuthorized) {
 					continue
@@ -350,21 +333,15 @@ func CheckAccountRolePrivilegesRevoked(t *testing.T) func(*terraform.State) erro
 // CheckDatabaseRolePrivilegesRevoked is a custom checks that should be later incorporated into generic CheckDestroy
 func CheckDatabaseRolePrivilegesRevoked(t *testing.T) func(*terraform.State) error {
 	t.Helper()
-	client := Client(t)
 
 	return func(s *terraform.State) error {
 		for _, rs := range s.RootModule().Resources {
 			if rs.Type != "snowflake_grant_privileges_to_database_role" {
 				continue
 			}
-			ctx := context.Background()
 
 			id := sdk.NewDatabaseObjectIdentifierFromFullyQualifiedName(rs.Primary.Attributes["database_role_name"])
-			grants, err := client.Grants.Show(ctx, &sdk.ShowGrantOptions{
-				To: &sdk.ShowGrantsTo{
-					DatabaseRole: id,
-				},
-			})
+			grants, err := TestClient().Grant.ShowGrantsToDatabaseRole(t, id)
 			if err != nil {
 				return err
 			}
@@ -386,23 +363,15 @@ func CheckDatabaseRolePrivilegesRevoked(t *testing.T) func(*terraform.State) err
 // CheckSharePrivilegesRevoked is a custom checks that should be later incorporated into generic CheckDestroy
 func CheckSharePrivilegesRevoked(t *testing.T) func(*terraform.State) error {
 	t.Helper()
-	client := Client(t)
 
 	return func(state *terraform.State) error {
 		for _, rs := range state.RootModule().Resources {
 			if rs.Type != "snowflake_grant_privileges_to_share" {
 				continue
 			}
-			ctx := context.Background()
 
 			id := sdk.NewExternalObjectIdentifierFromFullyQualifiedName(rs.Primary.Attributes["to_share"])
-			grants, err := client.Grants.Show(ctx, &sdk.ShowGrantOptions{
-				To: &sdk.ShowGrantsTo{
-					Share: &sdk.ShowGrantsToShare{
-						Name: sdk.NewAccountObjectIdentifier(id.Name()),
-					},
-				},
-			})
+			grants, err := TestClient().Grant.ShowGrantsToShare(t, sdk.NewAccountObjectIdentifier(id.Name()))
 			if err != nil {
 				return err
 			}
@@ -421,18 +390,12 @@ func CheckSharePrivilegesRevoked(t *testing.T) func(*terraform.State) error {
 // CheckUserPasswordPolicyAttachmentDestroy is a custom checks that should be later incorporated into generic CheckDestroy
 func CheckUserPasswordPolicyAttachmentDestroy(t *testing.T) func(*terraform.State) error {
 	t.Helper()
-	client := Client(t)
-
 	return func(s *terraform.State) error {
 		for _, rs := range s.RootModule().Resources {
 			if rs.Type != "snowflake_user_password_policy_attachment" {
 				continue
 			}
-			ctx := context.Background()
-			policyReferences, err := client.PolicyReferences.GetForEntity(ctx, sdk.NewGetForEntityPolicyReferenceRequest(
-				sdk.NewAccountObjectIdentifierFromFullyQualifiedName(rs.Primary.Attributes["user_name"]),
-				sdk.PolicyEntityDomainUser,
-			))
+			policyReferences, err := TestClient().PolicyReferences.GetPolicyReferences(t, sdk.NewAccountObjectIdentifierFromFullyQualifiedName(rs.Primary.Attributes["user_name"]), sdk.PolicyEntityDomainUser)
 			if err != nil {
 				if strings.Contains(err.Error(), "does not exist or not authorized") {
 					// Note: this can happen if the Policy Reference or the User has been deleted as well; in this case, ignore the error

diff --git a/pkg/acceptance/helpers/common.go b/pkg/acceptance/helpers/common.go
@@ -2,7 +2,6 @@ package helpers
 
 import (
 	"context"
-	"database/sql"
 	"fmt"
 	"log"
 	"testing"
@@ -71,20 +70,3 @@ func AssertErrorContainsPartsFunc(t *testing.T, parts []string) resource.ErrorCh
 		return nil
 	}
 }
-
-type PolicyReference struct {
-	PolicyDb          string         `db:"POLICY_DB"`
-	PolicySchema      string         `db:"POLICY_SCHEMA"`
-	PolicyName        string         `db:"POLICY_NAME"`
-	PolicyKind        string         `db:"POLICY_KIND"`
-	RefDatabaseName   string         `db:"REF_DATABASE_NAME"`
-	RefSchemaName     string         `db:"REF_SCHEMA_NAME"`
-	RefEntityName     string         `db:"REF_ENTITY_NAME"`
-	RefEntityDomain   string         `db:"REF_ENTITY_DOMAIN"`
-	RefColumnName     sql.NullString `db:"REF_COLUMN_NAME"`
-	RefArgColumnNames sql.NullString `db:"REF_ARG_COLUMN_NAMES"`
-	TagDatabase       sql.NullString `db:"TAG_DATABASE"`
-	TagSchema         sql.NullString `db:"TAG_SCHEMA"`
-	TagName           sql.NullString `db:"TAG_NAME"`
-	PolicyStatus      string         `db:"POLICY_STATUS"`
-}
diff --git a/pkg/acceptance/helpers/database_client.go b/pkg/acceptance/helpers/database_client.go
@@ -153,3 +153,34 @@ func (c *DatabaseClient) Describe(t *testing.T, id sdk.AccountObjectIdentifier)
 
 	return c.client().Describe(ctx, id)
 }
+
+// TODO [SNOW-1562172]: Create a better solution for this type of situations
+// We have to create test database from share before the actual test to check if the newly created share is ready
+// after previous test (there's some kind of issue or delay between cleaning up a share and creating a new one right after).
+func (c *DatabaseClient) CreateDatabaseFromShareTemporarily(t *testing.T, externalShareId sdk.ExternalObjectIdentifier) {
+	t.Helper()
+
+	databaseId := c.ids.RandomAccountObjectIdentifier()
+	err := c.client().CreateShared(context.Background(), databaseId, externalShareId, new(sdk.CreateSharedDatabaseOptions))
+	require.NoError(t, err)
+
+	require.Eventually(t, func() bool {
+		database, err := c.Show(t, databaseId)
+		if err != nil {
+			return false
+		}
+		// Origin is returned as "<revoked>" in those cases, because it's not valid sdk.ExternalObjectIdentifier parser sets it as nil.
+		// Once it turns into valid sdk.ExternalObjectIdentifier, we're ready to proceed with the actual test.
+		return database.Origin != nil
+	}, time.Minute, time.Second*6)
+
+	err = c.DropDatabase(t, databaseId)
+	require.NoError(t, err)
+}
+
+func (c *DatabaseClient) ShowAllReplicationDatabases(t *testing.T) ([]sdk.ReplicationDatabase, error) {
+	t.Helper()
+	ctx := context.Background()
+
+	return c.context.client.ReplicationFunctions.ShowReplicationDatabases(ctx, nil)
+}
diff --git a/pkg/acceptance/helpers/external_table_client.go b/pkg/acceptance/helpers/external_table_client.go
@@ -0,0 +1,34 @@
+package helpers
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
+	"github.com/stretchr/testify/require"
+)
+
+type ExternalTableClient struct {
+	context *TestClientContext
+	ids     *IdsGenerator
+}
+
+func NewExternalTableClient(context *TestClientContext, idsGenerator *IdsGenerator) *ExternalTableClient {
+	return &ExternalTableClient{
+		context: context,
+		ids:     idsGenerator,
+	}
+}
+
+func (c *ExternalTableClient) client() sdk.ExternalTables {
+	return c.context.client.ExternalTables
+}
+
+func (c *ExternalTableClient) PublishDataToStage(t *testing.T, stageId sdk.SchemaObjectIdentifier, data []byte) {
+	t.Helper()
+	ctx := context.Background()
+
+	_, err := c.context.client.ExecForTests(ctx, fmt.Sprintf(`copy into @%s/external_tables_test_data/test_data from (select parse_json('%s')) overwrite = true`, stageId.FullyQualifiedName(), string(data)))
+	require.NoError(t, err)
+}