From e600496a6e7ab81777155dffccb1dc474710d68d Mon Sep 17 00:00:00 2001
From: Artur Sawicki <artur.sawicki@snowflake.com>
Date: Tue, 23 Apr 2024 18:28:43 +0200
Subject: [PATCH] Use session policy helper in tests

---
 .../helpers/session_policy_client.go          |  4 ++--
 pkg/sdk/testint/accounts_integration_test.go  |  3 ++-
 pkg/sdk/testint/helpers_test.go               | 19 -------------------
 3 files changed, 4 insertions(+), 22 deletions(-)

diff --git a/pkg/acceptance/helpers/session_policy_client.go b/pkg/acceptance/helpers/session_policy_client.go
index 2b9ee97e8e..36949b9809 100644
--- a/pkg/acceptance/helpers/session_policy_client.go
+++ b/pkg/acceptance/helpers/session_policy_client.go
@@ -23,9 +23,9 @@ func (c *SessionPolicyClient) client() sdk.SessionPolicies {
 	return c.context.client.SessionPolicies
 }
 
-func (c *SessionPolicyClient) CreateSessionPolicy(t *testing.T, database *sdk.Database, schema *sdk.Schema) (*sdk.SessionPolicy, func()) {
+func (c *SessionPolicyClient) CreateSessionPolicy(t *testing.T) (*sdk.SessionPolicy, func()) {
 	t.Helper()
-	id := sdk.NewSchemaObjectIdentifier(database.Name, schema.Name, random.StringN(12))
+	id := sdk.NewSchemaObjectIdentifier(c.context.database, c.context.schema, random.StringN(12))
 	return c.CreateSessionPolicyWithOptions(t, id, sdk.NewCreateSessionPolicyRequest(id))
 }
 
diff --git a/pkg/sdk/testint/accounts_integration_test.go b/pkg/sdk/testint/accounts_integration_test.go
index d177dc53c0..e3f49568a6 100644
--- a/pkg/sdk/testint/accounts_integration_test.go
+++ b/pkg/sdk/testint/accounts_integration_test.go
@@ -248,8 +248,9 @@ func TestInt_AccountAlter(t *testing.T) {
 	})
 
 	t.Run("set and unset session policy", func(t *testing.T) {
-		sessionPolicyTest, sessionPolicyCleanup := createSessionPolicy(t, client, testDb(t), testSchema(t))
+		sessionPolicyTest, sessionPolicyCleanup := testClientHelper().SessionPolicy.CreateSessionPolicy(t)
 		t.Cleanup(sessionPolicyCleanup)
+
 		opts := &sdk.AlterAccountOptions{
 			Set: &sdk.AccountSet{
 				SessionPolicy: sessionPolicyTest.ID(),
diff --git a/pkg/sdk/testint/helpers_test.go b/pkg/sdk/testint/helpers_test.go
index 92163c0cbe..105b77e562 100644
--- a/pkg/sdk/testint/helpers_test.go
+++ b/pkg/sdk/testint/helpers_test.go
@@ -130,25 +130,6 @@ func createNetworkPolicy(t *testing.T, client *sdk.Client, req *sdk.CreateNetwor
 	}
 }
 
-func createSessionPolicy(t *testing.T, client *sdk.Client, database *sdk.Database, schema *sdk.Schema) (*sdk.SessionPolicy, func()) {
-	t.Helper()
-	id := sdk.NewSchemaObjectIdentifier(database.Name, schema.Name, random.StringN(12))
-	return createSessionPolicyWithOptions(t, client, id, sdk.NewCreateSessionPolicyRequest(id))
-}
-
-func createSessionPolicyWithOptions(t *testing.T, client *sdk.Client, id sdk.SchemaObjectIdentifier, request *sdk.CreateSessionPolicyRequest) (*sdk.SessionPolicy, func()) {
-	t.Helper()
-	ctx := context.Background()
-	err := client.SessionPolicies.Create(ctx, request)
-	require.NoError(t, err)
-	sessionPolicy, err := client.SessionPolicies.ShowByID(ctx, id)
-	require.NoError(t, err)
-	return sessionPolicy, func() {
-		err := client.SessionPolicies.Drop(ctx, sdk.NewDropSessionPolicyRequest(id))
-		require.NoError(t, err)
-	}
-}
-
 func createResourceMonitor(t *testing.T, client *sdk.Client) (*sdk.ResourceMonitor, func()) {
 	t.Helper()
 	return createResourceMonitorWithOptions(t, client, &sdk.CreateResourceMonitorOptions{