Fix query tests

pkg/services/query/query_test.go

@@ -2,6 +2,7 @@ package query_test
 
 import (
 	"context"
+	"encoding/json"
 	"net/http"
 	"testing"
 
@@ -9,21 +10,32 @@ import (
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana/pkg/api/dtos"
+	"github.com/grafana/grafana/pkg/bus"
 	"github.com/grafana/grafana/pkg/components/simplejson"
 	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/plugins"
+	acmock "github.com/grafana/grafana/pkg/services/accesscontrol/mock"
+	datasources "github.com/grafana/grafana/pkg/services/datasources/service"
+	"github.com/grafana/grafana/pkg/services/featuremgmt"
 	"github.com/grafana/grafana/pkg/services/query"
 	"github.com/grafana/grafana/pkg/services/secrets"
+	"github.com/grafana/grafana/pkg/services/secrets/kvstore"
+
 	"github.com/stretchr/testify/require"
 )
 
 func TestQueryData(t *testing.T) {
 	t.Run("it attaches custom headers to the request", func(t *testing.T) {
-		tc := setup()
+		tc := setup(t)
 		tc.dataSourceCache.ds.JsonData = simplejson.NewFromAny(map[string]interface{}{"httpHeaderName1": "foo", "httpHeaderName2": "bar"})
-		tc.secretService.decryptedJson = map[string]string{"httpHeaderValue1": "test-header", "httpHeaderValue2": "test-header2"}
 
-		_, err := tc.queryService.QueryData(context.Background(), nil, true, metricRequest(), false)
+		secureJsonData, err := json.Marshal(map[string]string{"httpHeaderValue1": "test-header", "httpHeaderValue2": "test-header2"})
+		require.NoError(t, err)
+
+		err = tc.secretStore.Set(context.Background(), tc.dataSourceCache.ds.OrgId, tc.dataSourceCache.ds.Name, "datasource", string(secureJsonData))
+		require.NoError(t, err)
+
+		_, err = tc.queryService.QueryData(context.Background(), nil, true, metricRequest(), false)
 		require.Nil(t, err)
 
 		require.Equal(t, map[string]string{"foo": "test-header", "bar": "test-header2"}, tc.pluginContext.req.Headers)
@@ -36,7 +48,7 @@ func TestQueryData(t *testing.T) {
 		}
 		token = token.WithExtra(map[string]interface{}{"id_token": "id-token"})
 
-		tc := setup()
+		tc := setup(t)
 		tc.oauthTokenService.passThruEnabled = true
 		tc.oauthTokenService.token = token
 
@@ -51,26 +63,28 @@ func TestQueryData(t *testing.T) {
 	})
 }
 
-func setup() *testContext {
+func setup(t *testing.T) *testContext {
 	pc := &fakePluginClient{}
-	sc := &fakeSecretsService{}
 	dc := &fakeDataSourceCache{ds: &models.DataSource{}}
 	tc := &fakeOAuthTokenService{}
 	rv := &fakePluginRequestValidator{}
 
+	ss := kvstore.SetupTestService(t)
+	ds := datasources.ProvideService(bus.New(), nil, ss, nil, featuremgmt.WithFeatures(), acmock.New(), acmock.NewPermissionsServicesMock())
+
 	return &testContext{
 		pluginContext:          pc,
-		secretService:          sc,
+		secretStore:            ss,
 		dataSourceCache:        dc,
 		oauthTokenService:      tc,
 		pluginRequestValidator: rv,
-		queryService:           query.ProvideService(nil, dc, nil, rv, sc, pc, tc),
+		queryService:           query.ProvideService(nil, dc, nil, rv, ds, pc, tc),
 	}
 }
 
 type testContext struct {
 	pluginContext          *fakePluginClient
-	secretService          *fakeSecretsService
+	secretStore            kvstore.SecretsKVStore
 	dataSourceCache        *fakeDataSourceCache
 	oauthTokenService      *fakeOAuthTokenService
 	pluginRequestValidator *fakePluginRequestValidator