fix test

astronomer · Jul 8, 2024 · 3333941 · 3333941
1 parent 7819870
commit 3333941
Showing 1 changed file with 18 additions and 4 deletions.
diff --git a/internal/provider/datasources/data_source_users_test.go b/internal/provider/datasources/data_source_users_test.go
@@ -25,19 +25,19 @@ func TestAcc_DataSourceUsers(t *testing.T) {
 			{
 				Config: astronomerprovider.ProviderConfig(t, true) + users(tfVarName),
 				Check: resource.ComposeTestCheckFunc(
-					checkUsers(tfVarName),
+					checkUsers(tfVarName, false, false),
 				),
 			},
 			{
 				Config: astronomerprovider.ProviderConfig(t, true) + usersFilter(tfVarName, "workspace_id", tfWorkspaceId),
 				Check: resource.ComposeTestCheckFunc(
-					checkUsers(tfVarName),
+					checkUsers(tfVarName, true, false),
 				),
 			},
 			{
 				Config: astronomerprovider.ProviderConfig(t, true) + usersFilter(tfVarName, "deployment_id", tfDeploymentId),
 				Check: resource.ComposeTestCheckFunc(
-					checkUsers(tfVarName),
+					checkUsers(tfVarName, false, true),
 				),
 			},
 		},
@@ -56,7 +56,7 @@ data astro_users "%v" {
 }`, tfVarName, filter, filterId)
 }
 
-func checkUsers(tfVarName string) resource.TestCheckFunc {
+func checkUsers(tfVarName string, filterWorkspaceId bool, filterDeploymentId bool) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
 		instanceState, numUsers, err := utils.GetDataSourcesLength(s, tfVarName, "users")
 		if err != nil {
@@ -93,6 +93,20 @@ func checkUsers(tfVarName string) resource.TestCheckFunc {
 		if instanceState.Attributes[organizationRole] == "" {
 			return fmt.Errorf("expected 'organization_role' to be set")
 		}
+		if filterWorkspaceId {
+			workspaceRoles := fmt.Sprintf("users.%d.workspace_roles", usersIdx)
+			if len(instanceState.Attributes[workspaceRoles]) == 0 {
+				fmt.Printf(instanceState.Attributes[workspaceRoles])
+				return fmt.Errorf("expected 'workspace_roles' to be set")
+			}
+		}
+		if filterDeploymentId {
+			deploymentRoles := fmt.Sprintf("users.%d.deployment_roles", usersIdx)
+			if len(instanceState.Attributes[deploymentRoles]) == 0 {
+				fmt.Printf(instanceState.Attributes[deploymentRoles])
+				return fmt.Errorf("expected 'deployment_roles' to be set")
+			}
+		}
 		createdAt := fmt.Sprintf("users.%d.created_at", usersIdx)
 		if instanceState.Attributes[createdAt] == "" {
 			return fmt.Errorf("expected 'created_at' to be set")