From 458fc6961677ff4c1d79fb98abdc7fafe32129de Mon Sep 17 00:00:00 2001
From: Aaron Godin <aaron.godin@grafana.com>
Date: Fri, 4 Oct 2024 06:35:25 -0500
Subject: [PATCH] IAM: Add test for AddDataSource managed permissions (#94113)

Add test for AddDataSource managed permissions and fix control flow issues
---
 .../datasources/service/datasource.go         | 12 ++++--
 .../datasources/service/datasource_test.go    | 38 ++++++++++++++++---
 2 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/pkg/services/datasources/service/datasource.go b/pkg/services/datasources/service/datasource.go
index ef51ebef430..495861cc941 100644
--- a/pkg/services/datasources/service/datasource.go
+++ b/pkg/services/datasources/service/datasource.go
@@ -255,7 +255,7 @@ func (s *Service) AddDataSource(ctx context.Context, cmd *datasources.AddDataSou
 	}
 
 	var dataSource *datasources.DataSource
-	return dataSource, s.db.InTransaction(ctx, func(ctx context.Context) error {
+	err = s.db.InTransaction(ctx, func(ctx context.Context) error {
 		var err error
 
 		cmd.EncryptedSecureJsonData = make(map[string][]byte)
@@ -293,12 +293,18 @@ func (s *Service) AddDataSource(ctx context.Context, cmd *datasources.AddDataSou
 			if cmd.UserID != 0 {
 				permissions = append(permissions, accesscontrol.SetResourcePermissionCommand{UserID: cmd.UserID, Permission: "Admin"})
 			}
-			_, err = s.permissionsService.SetPermissions(ctx, cmd.OrgID, dataSource.UID, permissions...)
-			return err
+			if _, err = s.permissionsService.SetPermissions(ctx, cmd.OrgID, dataSource.UID, permissions...); err != nil {
+				return err
+			}
 		}
 
 		return nil
 	})
+	if err != nil {
+		return nil, err
+	}
+
+	return dataSource, nil
 }
 
 // This will valid validate the instance settings return a version that is safe to be saved
diff --git a/pkg/services/datasources/service/datasource_test.go b/pkg/services/datasources/service/datasource_test.go
index 3433abdd2e2..c13456d4cf0 100644
--- a/pkg/services/datasources/service/datasource_test.go
+++ b/pkg/services/datasources/service/datasource_test.go
@@ -3,6 +3,7 @@ package service
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -104,6 +105,27 @@ func TestService_AddDataSource(t *testing.T) {
 		require.EqualError(t, err, "[datasource.urlInvalid] max length is 255")
 	})
 
+	t.Run("should fail if the datasource managed permissions fail", func(t *testing.T) {
+		dsService := initDSService(t)
+		enableRBACManagedPermissions(t, dsService.cfg)
+		dsService.permissionsService = &actest.FakePermissionsService{
+			ExpectedErr: errors.New("failed to set datasource permissions"),
+		}
+		dsService.pluginStore = &pluginstore.FakePluginStore{
+			PluginList: []pluginstore.Plugin{},
+		}
+
+		cmd := &datasources.AddDataSourceCommand{
+			OrgID: 1,
+			Type:  datasources.DS_TESTDATA,
+			Name:  "test",
+		}
+
+		ds, err := dsService.AddDataSource(context.Background(), cmd)
+		assert.Nil(t, ds)
+		assert.ErrorContains(t, err, "failed to set datasource permissions")
+	})
+
 	t.Run("if a plugin has an API version defined (EXPERIMENTAL)", func(t *testing.T) {
 		t.Run("should success to run admission hooks", func(t *testing.T) {
 			dsService := initDSService(t)
@@ -580,11 +602,8 @@ func TestService_DeleteDataSource(t *testing.T) {
 		permissionSvc := acmock.NewMockedPermissionsService()
 		permissionSvc.On("SetPermissions", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]accesscontrol.ResourcePermission{}, nil).Once()
 		permissionSvc.On("DeleteResourcePermissions", mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
-
-		f := ini.Empty()
-		f.Section("rbac").Key("resources_with_managed_permissions_on_creation").SetValue("datasource")
-		cfg, err := setting.NewCfgFromINIFile(f)
-		require.NoError(t, err)
+		cfg := &setting.Cfg{}
+		enableRBACManagedPermissions(t, cfg)
 		dsService, err := ProvideService(sqlStore, secretsService, secretsStore, cfg, featuremgmt.WithFeatures(), acmock.New(), permissionSvc, quotaService, &pluginstore.FakePluginStore{}, &pluginfakes.FakePluginClient{}, nil)
 		require.NoError(t, err)
 
@@ -1521,6 +1540,15 @@ func initDSService(t *testing.T) *Service {
 	return dsService
 }
 
+func enableRBACManagedPermissions(t testing.TB, cfg *setting.Cfg) {
+	t.Helper()
+	f := ini.Empty()
+	f.Section("rbac").Key("resources_with_managed_permissions_on_creation").SetValue("datasource")
+	tempCfg, err := setting.NewCfgFromINIFile(f)
+	cfg.RBAC = tempCfg.RBAC
+	require.NoError(t, err)
+}
+
 const caCert string = `-----BEGIN CERTIFICATE-----
 MIIDATCCAemgAwIBAgIJAMQ5hC3CPDTeMA0GCSqGSIb3DQEBCwUAMBcxFTATBgNV
 BAMMDGNhLWs4cy1zdGhsbTAeFw0xNjEwMjcwODQyMjdaFw00NDAzMTQwODQyMjda