Address review comments

Author: sfc-gh-jcieslak

pkg/acceptance/bettertestspoc/assert/objectassert/security_integration_snowflake_ext.go

@@ -2,20 +2,9 @@ package resourceshowoutputassert
 
 import (
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/bettertestspoc/assert"
-	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/testenvs"
 )
 
 func (s *SecurityIntegrationShowOutputAssert) HasCreatedOnNotEmpty() *SecurityIntegrationShowOutputAssert {
 	s.AddAssertion(assert.ValuePresent("show_output.0.created_on"))
 	return s
 }
-
-func (s *SecurityIntegrationShowOutputAssert) HasEnabledSnowflakeDefault() *SecurityIntegrationShowOutputAssert {
-	env := testenvs.GetSnowflakeEnvironmentWithProdDefault()
-	enabled := false
-	if env == testenvs.SnowflakeNonProdEnvironment {
-		enabled = true
-	}
-	s.AddAssertion(assert.ResourceShowOutputBoolValueSet("enabled", enabled))
-	return s
-}

pkg/acceptance/bettertestspoc/assert/resourceshowoutputassert/security_integration_show_output_ext.go

@@ -0,0 +1,8 @@
+package snowflakedefaults
+
+func EnabledValueForSnowflakeOauthSecurityIntegration() bool {
+	if getSnowflakeEnvironmentWithProdDefault() == SnowflakeNonProdEnvironment {
+		return true
+	}
+	return false
+}

pkg/acceptance/snowflakedefaults/security_integration.go

@@ -0,0 +1,45 @@
+package snowflakedefaults
+
+import (
+	"fmt"
+	"log"
+	"os"
+	"slices"
+	"strings"
+
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/testenvs"
+)
+
+type SnowflakeEnvironment string
+
+const (
+	SnowflakeProdEnvironment    SnowflakeEnvironment = "PROD"
+	SnowflakeNonProdEnvironment SnowflakeEnvironment = "NON_PROD"
+)
+
+var allSnowflakeEnvironments = []SnowflakeEnvironment{
+	SnowflakeProdEnvironment,
+	SnowflakeNonProdEnvironment,
+}
+
+func parseSnowflakeEnvironment(environment string) (SnowflakeEnvironment, error) {
+	snowflakeEnvironment := SnowflakeEnvironment(strings.ToUpper(environment))
+	if slices.Contains(allSnowflakeEnvironments, snowflakeEnvironment) {
+		return snowflakeEnvironment, nil
+	}
+	return "", fmt.Errorf("invalid Snowflake environment: %s, valid values are: %v", environment, allSnowflakeEnvironments)
+}
+
+func getSnowflakeEnvironmentWithProdDefault() SnowflakeEnvironment {
+	env := os.Getenv(string(testenvs.SnowflakeTestingEnvironment))
+	if env == "" {
+		log.Printf("[DEBUG] Snowflake environment variable %s not set, returning default PROD environment", testenvs.SnowflakeTestingEnvironment)
+		return SnowflakeProdEnvironment
+	}
+	snowflakeEnvironment, err := parseSnowflakeEnvironment(env)
+	if err != nil {
+		log.Printf("[DEBUG] Failed to parse Snowflake environment variable (%s), returning default PROD environment, err = %s", testenvs.SnowflakeTestingEnvironment, err)
+		return SnowflakeProdEnvironment
+	}
+	return snowflakeEnvironment
+}

pkg/acceptance/snowflakedefaults/snowflake_environment.go

@@ -0,0 +1,66 @@
+package snowflakedefaults
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseSnowflakeEnvironment(t *testing.T) {
+	t.Run("valid values", func(t *testing.T) {
+		testCases := map[string]struct {
+			input    string
+			expected SnowflakeEnvironment
+		}{
+			"exact PROD": {
+				input:    string(SnowflakeProdEnvironment),
+				expected: SnowflakeProdEnvironment,
+			},
+			"lowercase prod": {
+				input:    "prod",
+				expected: SnowflakeProdEnvironment,
+			},
+			"exact NON_PROD": {
+				input:    string(SnowflakeNonProdEnvironment),
+				expected: SnowflakeNonProdEnvironment,
+			},
+			"lowercase non_prod": {
+				input:    "non_prod",
+				expected: SnowflakeNonProdEnvironment,
+			},
+			"mixed case prod": {
+				input:    "PrOd",
+				expected: SnowflakeProdEnvironment,
+			},
+		}
+
+		for name, tc := range testCases {
+			tc := tc
+			t.Run(name, func(t *testing.T) {
+				actual, err := parseSnowflakeEnvironment(tc.input)
+				require.NoError(t, err)
+				require.Equal(t, tc.expected, actual)
+			})
+		}
+	})
+
+	t.Run("invalid values", func(t *testing.T) {
+		testCases := []string{
+			"",
+			"dev",
+			"prod-env",
+			"non prod",
+		}
+
+		for _, input := range testCases {
+			input := input
+			t.Run(input, func(t *testing.T) {
+				t.Parallel()
+
+				_, err := parseSnowflakeEnvironment(input)
+				require.Error(t, err)
+				require.ErrorContains(t, err, "invalid Snowflake environment")
+			})
+		}
+	})
+}

pkg/acceptance/snowflakedefaults/snowflake_environment_test.go

@@ -0,0 +1,14 @@
+package snowflakedefaults
+
+import (
+	"fmt"
+
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
+)
+
+func StageIdentifierOutputFormatForStreamOnDirectoryTable(id sdk.SchemaObjectIdentifier) string {
+	if getSnowflakeEnvironmentWithProdDefault() == SnowflakeNonProdEnvironment {
+		return fmt.Sprintf(`"%s"."%s".%s`, id.DatabaseName(), id.SchemaName(), id.Name())
+	}
+	return id.Name()
+}

pkg/acceptance/snowflakedefaults/stream_on_directory_table.go

@@ -5,9 +5,11 @@ package testacc
 import (
 	"fmt"
 	"regexp"
+	"strconv"
 	"testing"
 
 	accconfig "github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/bettertestspoc/config"
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/snowflakedefaults"
 	resourcenames "github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/resources"
 
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/bettertestspoc/assert"
@@ -17,7 +19,6 @@ import (
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/bettertestspoc/assert/resourceshowoutputassert"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/bettertestspoc/config/model"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/helpers/random"
-	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/testenvs"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/resources"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
 	"github.com/hashicorp/terraform-plugin-testing/helper/resource"
@@ -52,18 +53,15 @@ func TestAcc_OauthIntegrationForCustomClients_BasicUseCase(t *testing.T) {
 		WithPreAuthorizedRoles(preAuthorizedRole.ID()).
 		WithComment(comment)
 
-	snowflakeEnv := testenvs.GetSnowflakeEnvironmentWithProdDefault()
-	enabledSnowflakeDefault := resources.BooleanFalse
-	if snowflakeEnv == testenvs.SnowflakeNonProdEnvironment {
-		enabledSnowflakeDefault = resources.BooleanTrue
-	}
+	enabledSnowflakeDefault := snowflakedefaults.EnabledValueForSnowflakeOauthSecurityIntegration()
+	enabledSnowflakeDefaultString := strconv.FormatBool(enabledSnowflakeDefault)
 
 	assertBasic := []assert.TestCheckFuncProvider{
 		objectassert.SecurityIntegration(t, id).
 			HasName(id.Name()).
 			HasIntegrationType("OAUTH - CUSTOM").
 			HasCategory("SECURITY").
-			HasEnabledSnowflakeDefault().
+			HasEnabled(enabledSnowflakeDefault).
 			HasComment(""),
 
 		resourceassert.OauthIntegrationForCustomClientsResource(t, basic.ResourceReference()).
@@ -91,7 +89,7 @@ func TestAcc_OauthIntegrationForCustomClients_BasicUseCase(t *testing.T) {
 		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.#", "1")),
 		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.oauth_client_type.0.value", "CONFIDENTIAL")),
 		assert.Check(resource.TestCheckNoResourceAttr(basic.ResourceReference(), "describe_output.0.oauth_redirect_uri.0.value")),
-		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.enabled.0.value", enabledSnowflakeDefault)),
+		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.enabled.0.value", enabledSnowflakeDefaultString)),
 		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.oauth_allow_non_tls_redirect_uri.0.value", resources.BooleanFalse)),
 		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.oauth_enforce_pkce.0.value", resources.BooleanFalse)),
 		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.oauth_use_secondary_roles.0.value", "NONE")),

pkg/acceptance/testenvs/snowflake_environment.go

@@ -8,6 +8,7 @@ package testacc
 import (
 	"fmt"
 	"regexp"
+	"strconv"
 	"strings"
 	"testing"
 
@@ -25,7 +26,7 @@ import (
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/helpers/random"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/importchecks"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/planchecks"
-	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/testenvs"
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/snowflakedefaults"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/datasources"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/snowflakeroles"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/resources"
@@ -50,18 +51,15 @@ func TestAcc_OauthIntegrationForPartnerApplications_BasicUseCase(t *testing.T) {
 		WithOauthUseSecondaryRoles(string(sdk.OauthSecurityIntegrationUseSecondaryRolesImplicit)).
 		WithComment(comment)
 
-	snowflakeEnv := testenvs.GetSnowflakeEnvironmentWithProdDefault()
-	enabledSnowflakeDefault := r.BooleanFalse
-	if snowflakeEnv == testenvs.SnowflakeNonProdEnvironment {
-		enabledSnowflakeDefault = r.BooleanTrue
-	}
+	enabledSnowflakeDefault := snowflakedefaults.EnabledValueForSnowflakeOauthSecurityIntegration()
+	enabledSnowflakeDefaultString := strconv.FormatBool(enabledSnowflakeDefault)
 
 	assertBasic := []assert.TestCheckFuncProvider{
 		objectassert.SecurityIntegration(t, id).
 			HasName(id.Name()).
 			HasIntegrationType("OAUTH - LOOKER").
 			HasCategory("SECURITY").
-			HasEnabledSnowflakeDefault().
+			HasEnabled(enabledSnowflakeDefault).
 			HasComment(""),
 
 		resourceassert.OauthIntegrationForPartnerApplicationsResource(t, basic.ResourceReference()).
@@ -81,13 +79,13 @@ func TestAcc_OauthIntegrationForPartnerApplications_BasicUseCase(t *testing.T) {
 			HasName(id.Name()).
 			HasIntegrationType("OAUTH - LOOKER").
 			HasCategory("SECURITY").
-			HasEnabledSnowflakeDefault().
+			HasEnabled(enabledSnowflakeDefault).
 			HasComment(""),
 
 		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.#", "1")),
 		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.oauth_client_type.0.value", "CONFIDENTIAL")),
 		assert.Check(resource.TestCheckNoResourceAttr(basic.ResourceReference(), "describe_output.0.oauth_redirect_uri.0.value")),
-		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.enabled.0.value", enabledSnowflakeDefault)),
+		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.enabled.0.value", enabledSnowflakeDefaultString)),
 		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.oauth_use_secondary_roles.0.value", "NONE")),
 		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.blocked_roles_list.0.value", "ACCOUNTADMIN,SECURITYADMIN")),
 		assert.Check(resource.TestCheckResourceAttr(basic.ResourceReference(), "describe_output.0.oauth_issue_refresh_tokens.0.value", r.BooleanTrue)),

pkg/testacc/resource_oauth_integration_for_custom_clients_acceptance_test.go

@@ -16,7 +16,7 @@ import (
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/bettertestspoc/config"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/bettertestspoc/config/model"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/helpers/random"
-	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/testenvs"
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/snowflakedefaults"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/resources"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
 	"github.com/hashicorp/terraform-plugin-testing/helper/resource"
@@ -37,11 +37,7 @@ func TestAcc_StreamOnDirectoryTable_BasicUseCase(t *testing.T) {
 	complete := model.StreamOnDirectoryTable("test", id.DatabaseName(), id.SchemaName(), id.Name(), stage.ID().FullyQualifiedName()).
 		WithComment(comment)
 
-	snowflakeEnv := testenvs.GetSnowflakeEnvironmentWithProdDefault()
-	expectedStageId := stage.ID().Name()
-	if snowflakeEnv == testenvs.SnowflakeNonProdEnvironment {
-		expectedStageId = fmt.Sprintf(`"%s"."%s".%s`, stage.DatabaseName, stage.SchemaName, stage.Name)
-	}
+	expectedStageId := snowflakedefaults.StageIdentifierOutputFormatForStreamOnDirectoryTable(stage.ID())
 
 	assertBasic := []assert.TestCheckFuncProvider{
 		resourceassert.StreamOnDirectoryTableResource(t, basic.ResourceReference()).