grafana/pkg/services/query/query_test.go

package query

import (
	"bytes"
	"context"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"

	"github.com/grafana/grafana-plugin-sdk-go/backend"
	"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/oauth2"

	"github.com/grafana/grafana/pkg/api/dtos"
	"github.com/grafana/grafana/pkg/components/simplejson"
	"github.com/grafana/grafana/pkg/expr"
	"github.com/grafana/grafana/pkg/infra/db"
	"github.com/grafana/grafana/pkg/infra/httpclient/httpclientprovider"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/plugins"
	acmock "github.com/grafana/grafana/pkg/services/accesscontrol/mock"
	"github.com/grafana/grafana/pkg/services/datasources"
	fakeDatasources "github.com/grafana/grafana/pkg/services/datasources/fakes"
	dsSvc "github.com/grafana/grafana/pkg/services/datasources/service"
	"github.com/grafana/grafana/pkg/services/featuremgmt"
	"github.com/grafana/grafana/pkg/services/quota/quotatest"
	"github.com/grafana/grafana/pkg/services/secrets/fakes"
	secretskvs "github.com/grafana/grafana/pkg/services/secrets/kvstore"
	secretsmng "github.com/grafana/grafana/pkg/services/secrets/manager"
	"github.com/grafana/grafana/pkg/services/user"
	"github.com/grafana/grafana/pkg/setting"
)

func TestParseMetricRequest(t *testing.T) {
	t.Run("Test a simple single datasource query", func(t *testing.T) {
		tc := setup(t)
		json, err := simplejson.NewJson([]byte(`{
			"keepCookies": [ "cookie1", "cookie3", "login" ]
		}`))
		require.NoError(t, err)
		tc.dataSourceCache.dsByUid = func(ctx context.Context, datasourceUID string, user *user.SignedInUser, skipCache bool) (*datasources.DataSource, error) {
			if datasourceUID == "gIEkMvIVz" {
				return &datasources.DataSource{
					Uid:      "gIEkMvIVz",
					JsonData: json,
				}, nil
			}

			return nil, nil
		}

		token := &oauth2.Token{
			TokenType:   "bearer",
			AccessToken: "access-token",
		}
		token = token.WithExtra(map[string]interface{}{"id_token": "id-token"})

		tc.oauthTokenService.passThruEnabled = true
		tc.oauthTokenService.token = token

		mr := metricRequestWithQueries(t, `{
			"refId": "A",
			"datasource": {
				"uid": "gIEkMvIVz",
				"type": "postgres"
			}
		}`, `{
			"refId": "B",
			"datasource": {
				"uid": "gIEkMvIVz",
				"type": "postgres"
			}
		}`)
		parsedReq, err := tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
		require.NoError(t, err)
		require.NotNil(t, parsedReq)
		assert.False(t, parsedReq.hasExpression)
		assert.Len(t, parsedReq.parsedQueries, 1)
		assert.Contains(t, parsedReq.parsedQueries, "gIEkMvIVz")
		assert.Len(t, parsedReq.getFlattenedQueries(), 2)

		t.Run("createDataSourceQueryEnrichers should return 0 enrichers when no HTTP request", func(t *testing.T) {
			enrichers := parsedReq.createDataSourceQueryEnrichers(context.Background(), nil, tc.oauthTokenService, []string{})
			require.Empty(t, enrichers)
		})

		t.Run("createDataSourceQueryEnrichers should return 1 enricher", func(t *testing.T) {
			parsedReq.httpRequest = httptest.NewRequest(http.MethodGet, "/", nil)
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "cookie1"})
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "cookie2"})
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "cookie3"})
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "login"})

			enrichers := parsedReq.createDataSourceQueryEnrichers(context.Background(), nil, tc.oauthTokenService, []string{"login"})
			require.Len(t, enrichers, 1)
			require.NotNil(t, enrichers["gIEkMvIVz"])
			req := &backend.QueryDataRequest{}
			ctx := enrichers["gIEkMvIVz"](context.Background(), req)
			require.Len(t, req.Headers, 3)
			require.Equal(t, "Bearer access-token", req.Headers["Authorization"])
			require.Equal(t, "id-token", req.Headers["X-ID-Token"])
			require.Equal(t, "cookie1=; cookie3=", req.Headers["Cookie"])
			middlewares := httpclient.ContextualMiddlewareFromContext(ctx)
			require.Len(t, middlewares, 2)
			require.Equal(t, httpclientprovider.ForwardedCookiesMiddlewareName, middlewares[0].(httpclient.MiddlewareName).MiddlewareName())
			require.Equal(t, httpclientprovider.ForwardedOAuthIdentityMiddlewareName, middlewares[1].(httpclient.MiddlewareName).MiddlewareName())
		})
	})

	t.Run("Test a single datasource query with expressions", func(t *testing.T) {
		tc := setup(t)
		json, err := simplejson.NewJson([]byte(`{
			"keepCookies": [ "cookie1", "cookie3", "login" ]
		}`))
		require.NoError(t, err)
		tc.dataSourceCache.dsByUid = func(ctx context.Context, datasourceUID string, user *user.SignedInUser, skipCache bool) (*datasources.DataSource, error) {
			if datasourceUID == "gIEkMvIVz" {
				return &datasources.DataSource{
					Uid:      "gIEkMvIVz",
					JsonData: json,
				}, nil
			}

			return nil, nil
		}

		token := &oauth2.Token{
			TokenType:   "bearer",
			AccessToken: "access-token",
		}
		token = token.WithExtra(map[string]interface{}{"id_token": "id-token"})

		tc.oauthTokenService.passThruEnabled = true
		tc.oauthTokenService.token = token

		mr := metricRequestWithQueries(t, `{
			"refId": "A",
			"datasource": {
				"uid": "gIEkMvIVz",
				"type": "postgres"
			}
		}`, `{
			"refId": "B",
			"datasource": {
				"type": "__expr__",
				"uid": "__expr__",
				"name": "Expression"
			},
			"type": "math",
			"expression": "$A - 50"
		}`)
		parsedReq, err := tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
		require.NoError(t, err)
		require.NotNil(t, parsedReq)
		assert.True(t, parsedReq.hasExpression)
		assert.Len(t, parsedReq.parsedQueries, 2)
		assert.Contains(t, parsedReq.parsedQueries, "gIEkMvIVz")
		assert.Len(t, parsedReq.getFlattenedQueries(), 2)
		// Make sure we end up with something valid
		_, err = tc.queryService.handleExpressions(context.Background(), tc.signedInUser, parsedReq)
		assert.NoError(t, err)

		t.Run("createDataSourceQueryEnrichers should return 1 enricher", func(t *testing.T) {
			parsedReq.httpRequest = httptest.NewRequest(http.MethodGet, "/", nil)
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "cookie1"})
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "cookie2"})
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "cookie3"})
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "login"})

			enrichers := parsedReq.createDataSourceQueryEnrichers(context.Background(), nil, tc.oauthTokenService, []string{"login"})
			require.Len(t, enrichers, 1)
			require.NotNil(t, enrichers["gIEkMvIVz"])

			req := &backend.QueryDataRequest{}
			ctx := enrichers["gIEkMvIVz"](context.Background(), req)
			require.Len(t, req.Headers, 3)
			require.Equal(t, "Bearer access-token", req.Headers["Authorization"])
			require.Equal(t, "id-token", req.Headers["X-ID-Token"])
			require.Equal(t, "cookie1=; cookie3=", req.Headers["Cookie"])
			middlewares := httpclient.ContextualMiddlewareFromContext(ctx)
			require.Len(t, middlewares, 2)
			require.Equal(t, httpclientprovider.ForwardedCookiesMiddlewareName, middlewares[0].(httpclient.MiddlewareName).MiddlewareName())
			require.Equal(t, httpclientprovider.ForwardedOAuthIdentityMiddlewareName, middlewares[1].(httpclient.MiddlewareName).MiddlewareName())
		})
	})

	t.Run("Test a simple mixed datasource query", func(t *testing.T) {
		tc := setup(t)
		json, err := simplejson.NewJson([]byte(`{
			"keepCookies": [ "cookie1", "cookie3", "login" ]
		}`))
		require.NoError(t, err)
		json2, err := simplejson.NewJson([]byte(`{
			"keepCookies": [ "cookie2" ]
		}`))
		require.NoError(t, err)
		tc.dataSourceCache.dsByUid = func(ctx context.Context, datasourceUID string, user *user.SignedInUser, skipCache bool) (*datasources.DataSource, error) {
			if datasourceUID == "gIEkMvIVz" {
				return &datasources.DataSource{
					Uid:      "gIEkMvIVz",
					JsonData: json,
				}, nil
			}

			if datasourceUID == "sEx6ZvSVk" {
				return &datasources.DataSource{
					Uid:      "sEx6ZvSVk",
					JsonData: json2,
				}, nil
			}

			return nil, nil
		}

		token := &oauth2.Token{
			TokenType:   "bearer",
			AccessToken: "access-token",
		}
		token = token.WithExtra(map[string]interface{}{"id_token": "id-token"})

		tc.oauthTokenService.passThruEnabled = true
		tc.oauthTokenService.token = token

		mr := metricRequestWithQueries(t, `{
			"refId": "A",
			"datasource": {
				"uid": "gIEkMvIVz",
				"type": "postgres"
			}
		}`, `{
			"refId": "B",
			"datasource": {
				"uid": "sEx6ZvSVk",
				"type": "testdata"
			}
		}`, `{
			"refId": "C",
			"datasource": {
				"uid": "sEx6ZvSVk",
				"type": "testdata"
			}
		}`)
		parsedReq, err := tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
		require.NoError(t, err)
		require.NotNil(t, parsedReq)
		assert.False(t, parsedReq.hasExpression)
		assert.Len(t, parsedReq.parsedQueries, 2)
		assert.Contains(t, parsedReq.parsedQueries, "gIEkMvIVz")
		assert.Len(t, parsedReq.parsedQueries["gIEkMvIVz"], 1)
		assert.Contains(t, parsedReq.parsedQueries, "sEx6ZvSVk")
		assert.Len(t, parsedReq.parsedQueries["sEx6ZvSVk"], 2)
		assert.Len(t, parsedReq.getFlattenedQueries(), 3)

		t.Run("createDataSourceQueryEnrichers should return 2 enrichers", func(t *testing.T) {
			parsedReq.httpRequest = httptest.NewRequest(http.MethodGet, "/", nil)
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "cookie1"})
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "cookie2"})
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "cookie3"})
			parsedReq.httpRequest.AddCookie(&http.Cookie{Name: "login"})

			enrichers := parsedReq.createDataSourceQueryEnrichers(context.Background(), nil, tc.oauthTokenService, []string{"login"})
			require.Len(t, enrichers, 2)

			enricherOne := enrichers["gIEkMvIVz"]
			require.NotNil(t, enricherOne)
			reqOne := &backend.QueryDataRequest{}
			ctx := enricherOne(context.Background(), reqOne)
			require.Len(t, reqOne.Headers, 3)
			require.Equal(t, "Bearer access-token", reqOne.Headers["Authorization"])
			require.Equal(t, "id-token", reqOne.Headers["X-ID-Token"])
			require.Equal(t, "cookie1=; cookie3=", reqOne.Headers["Cookie"])
			middlewaresOne := httpclient.ContextualMiddlewareFromContext(ctx)
			require.Len(t, middlewaresOne, 2)
			require.Equal(t, httpclientprovider.ForwardedCookiesMiddlewareName, middlewaresOne[0].(httpclient.MiddlewareName).MiddlewareName())
			require.Equal(t, httpclientprovider.ForwardedOAuthIdentityMiddlewareName, middlewaresOne[1].(httpclient.MiddlewareName).MiddlewareName())

			enricherTwo := enrichers["sEx6ZvSVk"]
			require.NotNil(t, enricherTwo)
			reqTwo := &backend.QueryDataRequest{}
			ctx = enricherTwo(context.Background(), reqTwo)
			require.Len(t, reqTwo.Headers, 3)
			require.Equal(t, "Bearer access-token", reqTwo.Headers["Authorization"])
			require.Equal(t, "id-token", reqTwo.Headers["X-ID-Token"])
			require.Equal(t, "cookie2=", reqTwo.Headers["Cookie"])
			middlewaresTwo := httpclient.ContextualMiddlewareFromContext(ctx)
			require.Len(t, middlewaresTwo, 2)
			require.Equal(t, httpclientprovider.ForwardedCookiesMiddlewareName, middlewaresTwo[0].(httpclient.MiddlewareName).MiddlewareName())
			require.Equal(t, httpclientprovider.ForwardedOAuthIdentityMiddlewareName, middlewaresTwo[1].(httpclient.MiddlewareName).MiddlewareName())
		})
	})

	t.Run("Test a mixed datasource query with expressions", func(t *testing.T) {
		tc := setup(t)
		mr := metricRequestWithQueries(t, `{
			"refId": "A",
			"datasource": {
				"uid": "gIEkMvIVz",
				"type": "postgres"
			}
		}`, `{
			"refId": "B",
			"datasource": {
				"uid": "sEx6ZvSVk",
				"type": "testdata"
			}
		}`, `{
			"refId": "A_resample",
			"datasource": {
				"type": "__expr__",
				"uid": "__expr__",
				"name": "Expression"
			},
			"expression": "A",
			"type": "resample",
			"downsampler": "mean",
			"upsampler": "fillna",
			"window": "10s"
		}`, `{
			"refId": "B_resample",
			"datasource": {
				"type": "__expr__",
				"uid": "__expr__",
				"name": "Expression"
			},
			"expression": "B",
			"type": "resample",
			"downsampler": "mean",
			"upsampler": "fillna",
			"window": "10s"
		}`, `{
			"refId": "C",
			"datasource": {
				"type": "__expr__",
				"uid": "__expr__",
				"name": "Expression"
			},
			"type": "math",
			"expression": "$A_resample + $B_resample"
		}`)
		parsedReq, err := tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
		require.NoError(t, err)
		require.NotNil(t, parsedReq)
		assert.True(t, parsedReq.hasExpression)
		assert.Len(t, parsedReq.parsedQueries, 3)
		assert.Contains(t, parsedReq.parsedQueries, "gIEkMvIVz")
		assert.Contains(t, parsedReq.parsedQueries, "sEx6ZvSVk")
		assert.Len(t, parsedReq.getFlattenedQueries(), 5)
		// Make sure we end up with something valid
		_, err = tc.queryService.handleExpressions(context.Background(), tc.signedInUser, parsedReq)
		assert.NoError(t, err)

		t.Run("createDataSourceQueryEnrichers should return 2 enrichers", func(t *testing.T) {
			parsedReq.httpRequest = &http.Request{}
			enrichers := parsedReq.createDataSourceQueryEnrichers(context.Background(), nil, tc.oauthTokenService, []string{})
			require.Len(t, enrichers, 2)
			require.NotNil(t, enrichers["gIEkMvIVz"])
			require.NotNil(t, enrichers["sEx6ZvSVk"])
		})
	})

	t.Run("Header validation", func(t *testing.T) {
		tc := setup(t)
		mr := metricRequestWithQueries(t, `{
			"refId": "A",
			"datasource": {
				"uid": "gIEkMvIVz",
				"type": "postgres"
			}
		}`, `{
			"refId": "B",
			"datasource": {
				"uid": "sEx6ZvSVk",
				"type": "testdata"
			}
		}`)
		httpreq, _ := http.NewRequest(http.MethodPost, "http://localhost/", bytes.NewReader([]byte{}))
		httpreq.Header.Add("X-Datasource-Uid", "gIEkMvIVz")
		mr.HTTPRequest = httpreq
		_, err := tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
		require.NoError(t, err)

		// With the second value it is OK
		httpreq.Header.Add("X-Datasource-Uid", "sEx6ZvSVk")
		mr.HTTPRequest = httpreq
		_, err = tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
		require.NoError(t, err)

		// Single header with comma syntax
		httpreq, _ = http.NewRequest(http.MethodPost, "http://localhost/", bytes.NewReader([]byte{}))
		httpreq.Header.Set("X-Datasource-Uid", "gIEkMvIVz, sEx6ZvSVk")
		mr.HTTPRequest = httpreq
		_, err = tc.queryService.parseMetricRequest(context.Background(), tc.signedInUser, true, mr)
		require.NoError(t, err)
	})
}

func TestQueryDataMultipleSources(t *testing.T) {
	t.Run("can query multiple datasources", func(t *testing.T) {
		tc := setup(t)
		query1, err := simplejson.NewJson([]byte(`
			{
				"datasource": {
					"type": "mysql",
					"uid": "ds1"
				},
				"refId": "A"
			}
		`))
		require.NoError(t, err)
		query2, err := simplejson.NewJson([]byte(`
			{
				"datasource": {
					"type": "mysql",
					"uid": "ds2"
				},
				"refId": "B"
			}
		`))
		require.NoError(t, err)
		queries := []*simplejson.Json{query1, query2}
		reqDTO := dtos.MetricRequest{
			From:                       "2022-01-01",
			To:                         "2022-01-02",
			Queries:                    queries,
			Debug:                      false,
			PublicDashboardAccessToken: "abc123",
			HTTPRequest:                nil,
		}

		_, err = tc.queryService.QueryData(context.Background(), tc.signedInUser, true, reqDTO)

		require.NoError(t, err)
	})

	t.Run("can query multiple datasources with an expression present", func(t *testing.T) {
		tc := setup(t)
		query1, err := simplejson.NewJson([]byte(`
			{
				"datasource": {
					"type": "mysql",
					"uid": "ds1"
				},
				"refId": "A"
			}
		`))
		require.NoError(t, err)
		query2, err := simplejson.NewJson([]byte(`
			{
				"datasource": {
					"type": "mysql",
					"uid": "ds2"
				},
				"refId": "B"
			}
		`))
		require.NoError(t, err)
		query3, err := simplejson.NewJson([]byte(`
			{
				"datasource": {
					"name": "Expression",
					"type": "__expr__",
					"uid": "__expr__"
				},
				"expression": "$A + 1",
				"hide": false,
				"refId": "EXPRESSION",
				"type": "math"
			}
		`))
		require.NoError(t, err)
		queries := []*simplejson.Json{query1, query2, query3}
		reqDTO := dtos.MetricRequest{
			From:                       "2022-01-01",
			To:                         "2022-01-02",
			Queries:                    queries,
			Debug:                      false,
			PublicDashboardAccessToken: "abc123",
			HTTPRequest:                nil,
		}

		// without query parameter
		_, err = tc.queryService.QueryData(context.Background(), tc.signedInUser, true, reqDTO)
		require.NoError(t, err)

		httpreq, _ := http.NewRequest(http.MethodPost, "http://localhost/ds/query?expression=true", bytes.NewReader([]byte{}))
		httpreq.Header.Add("X-Datasource-Uid", "gIEkMvIVz")
		reqDTO.HTTPRequest = httpreq

		// with query parameter
		_, err = tc.queryService.QueryData(context.Background(), tc.signedInUser, true, reqDTO)
		require.NoError(t, err)
	})

	t.Run("error is returned in query when one of the queries fails", func(t *testing.T) {
		tc := setup(t)

		query1, _ := simplejson.NewJson([]byte(`
			{
				"datasource": {
					"type": "mysql",
					"uid": "ds1"
				},
				"refId": "A"
			}
		`))
		query2, _ := simplejson.NewJson([]byte(`
			{
				"datasource": {
					"type": "prometheus",
					"uid": "ds2"
				},
				"refId": "B",
				"queryType": "FAIL"
			}
		`))

		queries := []*simplejson.Json{query1, query2}

		reqDTO := dtos.MetricRequest{
			From:                       "2022-01-01",
			To:                         "2022-01-02",
			Queries:                    queries,
			Debug:                      false,
			PublicDashboardAccessToken: "abc123",
			HTTPRequest:                nil,
		}

		res, err := tc.queryService.QueryData(context.Background(), tc.signedInUser, true, reqDTO)

		require.NoError(t, err)
		require.Error(t, res.Responses["B"].Error)
		// Responses aren't mocked, so a "healthy" query will just return an empty response
		require.NotContains(t, res.Responses, "A")
	})
}

func TestQueryData(t *testing.T) {
	t.Run("it auth custom headers to the request", func(t *testing.T) {
		token := &oauth2.Token{
			TokenType:   "bearer",
			AccessToken: "access-token",
		}
		token = token.WithExtra(map[string]interface{}{"id_token": "id-token"})

		tc := setup(t)
		tc.oauthTokenService.passThruEnabled = true
		tc.oauthTokenService.token = token

		metricReq := metricRequest()
		httpReq, err := http.NewRequest(http.MethodGet, "/", nil)
		require.NoError(t, err)
		metricReq.HTTPRequest = httpReq

		_, err = tc.queryService.QueryData(context.Background(), nil, true, metricReq)
		require.Nil(t, err)

		expected := map[string]string{
			"Authorization": "Bearer access-token",
			"X-ID-Token":    "id-token",
		}
		require.Equal(t, expected, tc.pluginContext.req.Headers)
	})

	t.Run("it doesn't add cookie header to the request when keepCookies configured and no cookies provided", func(t *testing.T) {
		tc := setup(t)
		json, err := simplejson.NewJson([]byte(`{"keepCookies": [ "foo", "bar" ]}`))
		require.NoError(t, err)
		tc.dataSourceCache.ds.JsonData = json

		metricReq := metricRequest()
		httpReq, err := http.NewRequest(http.MethodGet, "/", nil)
		require.NoError(t, err)
		metricReq.HTTPRequest = httpReq
		_, err = tc.queryService.QueryData(context.Background(), tc.signedInUser, true, metricReq)
		require.NoError(t, err)

		require.Empty(t, tc.pluginContext.req.Headers)
	})

	t.Run("it adds cookie header to the request when keepCookies configured and cookie provided", func(t *testing.T) {
		tc := setup(t)
		json, err := simplejson.NewJson([]byte(`{"keepCookies": [ "foo", "bar" ]}`))
		require.NoError(t, err)
		tc.dataSourceCache.ds.JsonData = json

		metricReq := metricRequest()
		httpReq, err := http.NewRequest(http.MethodGet, "/", nil)
		require.NoError(t, err)
		httpReq.AddCookie(&http.Cookie{Name: "a"})
		httpReq.AddCookie(&http.Cookie{Name: "bar", Value: "rab"})
		httpReq.AddCookie(&http.Cookie{Name: "b"})
		httpReq.AddCookie(&http.Cookie{Name: "foo", Value: "oof"})
		httpReq.AddCookie(&http.Cookie{Name: "c"})
		metricReq.HTTPRequest = httpReq
		_, err = tc.queryService.QueryData(context.Background(), tc.signedInUser, true, metricReq)
		require.NoError(t, err)

		require.Equal(t, map[string]string{"Cookie": "bar=rab; foo=oof"}, tc.pluginContext.req.Headers)
	})

	t.Run("it doesn't adds cookie header to the request when keepCookies configured with login cookie name", func(t *testing.T) {
		tc := setup(t)
		tc.queryService.cfg.LoginCookieName = "grafana_session"
		json, err := simplejson.NewJson([]byte(`{"keepCookies": [ "grafana_session", "bar" ]}`))
		require.NoError(t, err)
		tc.dataSourceCache.ds.JsonData = json

		metricReq := metricRequest()
		httpReq, err := http.NewRequest(http.MethodGet, "/", nil)
		require.NoError(t, err)
		httpReq.AddCookie(&http.Cookie{Name: "a"})
		httpReq.AddCookie(&http.Cookie{Name: "bar", Value: "rab"})
		httpReq.AddCookie(&http.Cookie{Name: "b"})
		httpReq.AddCookie(&http.Cookie{Name: "foo", Value: "oof"})
		httpReq.AddCookie(&http.Cookie{Name: "c"})
		httpReq.AddCookie(&http.Cookie{Name: tc.queryService.cfg.LoginCookieName, Value: "val"})
		metricReq.HTTPRequest = httpReq
		_, err = tc.queryService.QueryData(context.Background(), tc.signedInUser, true, metricReq)
		require.NoError(t, err)

		require.Equal(t, map[string]string{"Cookie": "bar=rab"}, tc.pluginContext.req.Headers)
	})
}

func setup(t *testing.T) *testContext {
	t.Helper()
	pc := &fakePluginClient{}
	dc := &fakeDataSourceCache{ds: &datasources.DataSource{}}
	tc := &fakeOAuthTokenService{}
	rv := &fakePluginRequestValidator{}

	sqlStore := db.InitTestDB(t)
	secretsService := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
	ss := secretskvs.NewSQLSecretsKVStore(sqlStore, secretsService, log.New("test.logger"))
	ssvc := secretsmng.SetupTestService(t, fakes.NewFakeSecretsStore())
	quotaService := quotatest.New(false, nil)
	ds, err := dsSvc.ProvideService(nil, ssvc, ss, nil, featuremgmt.WithFeatures(), acmock.New(), acmock.NewMockedPermissionsService(), quotaService)
	require.NoError(t, err)
	fakeDatasourceService := &fakeDatasources.FakeDataSourceService{
		DataSources:           nil,
		SimulatePluginFailure: false,
	}
	exprService := expr.ProvideService(&setting.Cfg{ExpressionsEnabled: true}, pc, fakeDatasourceService)
	queryService := ProvideService(setting.NewCfg(), dc, exprService, rv, ds, pc, tc) // provider belonging to this package
	return &testContext{
		pluginContext:          pc,
		secretStore:            ss,
		dataSourceCache:        dc,
		oauthTokenService:      tc,
		pluginRequestValidator: rv,
		queryService:           queryService,
		signedInUser:           &user.SignedInUser{OrgID: 1},
	}
}

type testContext struct {
	pluginContext          *fakePluginClient
	secretStore            secretskvs.SecretsKVStore
	dataSourceCache        *fakeDataSourceCache
	oauthTokenService      *fakeOAuthTokenService
	pluginRequestValidator *fakePluginRequestValidator
	queryService           *Service // implementation belonging to this package
	signedInUser           *user.SignedInUser
}

func metricRequest() dtos.MetricRequest {
	q, _ := simplejson.NewJson([]byte(`{"datasourceId":1}`))
	return dtos.MetricRequest{
		From:    "",
		To:      "",
		Queries: []*simplejson.Json{q},
		Debug:   false,
	}
}

func metricRequestWithQueries(t *testing.T, rawQueries ...string) dtos.MetricRequest {
	t.Helper()
	queries := make([]*simplejson.Json, 0)
	for _, q := range rawQueries {
		json, err := simplejson.NewJson([]byte(q))
		require.NoError(t, err)
		queries = append(queries, json)
	}
	return dtos.MetricRequest{
		From:    "now-1h",
		To:      "now",
		Queries: queries,
		Debug:   false,
	}
}

type fakePluginRequestValidator struct {
	err error
}

func (rv *fakePluginRequestValidator) Validate(dsURL string, req *http.Request) error {
	return rv.err
}

type fakeOAuthTokenService struct {
	passThruEnabled bool
	token           *oauth2.Token
}

func (ts *fakeOAuthTokenService) GetCurrentOAuthToken(context.Context, *user.SignedInUser) *oauth2.Token {
	return ts.token
}

func (ts *fakeOAuthTokenService) IsOAuthPassThruEnabled(*datasources.DataSource) bool {
	return ts.passThruEnabled
}

func (ts *fakeOAuthTokenService) HasOAuthEntry(context.Context, *user.SignedInUser) (*models.UserAuth, bool, error) {
	return nil, false, nil
}

func (ts *fakeOAuthTokenService) TryTokenRefresh(context.Context, *models.UserAuth) error {
	return nil
}

func (ts *fakeOAuthTokenService) InvalidateOAuthTokens(context.Context, *models.UserAuth) error {
	return nil
}

type fakeDataSourceCache struct {
	ds      *datasources.DataSource
	dsByUid func(ctx context.Context, datasourceUID string, user *user.SignedInUser, skipCache bool) (*datasources.DataSource, error)
}

func (c *fakeDataSourceCache) GetDatasource(ctx context.Context, datasourceID int64, user *user.SignedInUser, skipCache bool) (*datasources.DataSource, error) {
	return c.ds, nil
}

func (c *fakeDataSourceCache) GetDatasourceByUID(ctx context.Context, datasourceUID string, user *user.SignedInUser, skipCache bool) (*datasources.DataSource, error) {
	if c.dsByUid != nil {
		return c.dsByUid(ctx, datasourceUID, user, skipCache)
	}

	return &datasources.DataSource{
		Uid: datasourceUID,
	}, nil
}

type fakePluginClient struct {
	plugins.Client
	req *backend.QueryDataRequest
}

func (c *fakePluginClient) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
	c.req = req

	// If an expression query ends up getting directly queried, we want it to return an error in our test.
	if req.PluginContext.PluginID == "__expr__" {
		return nil, errors.New("cant query an expression datasource")
	}

	if req.Queries[0].QueryType == "FAIL" {
		return nil, errors.New("plugin client failed")
	}

	return &backend.QueryDataResponse{Responses: make(backend.Responses)}, nil
}