grafana/pkg/services/login/authinfoimpl/store_test.go

package authinfoimpl

import (
	"context"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/oauth2"

	"github.com/grafana/grafana/pkg/infra/db"
	"github.com/grafana/grafana/pkg/services/login"
	secretstest "github.com/grafana/grafana/pkg/services/secrets/fakes"
	"github.com/grafana/grafana/pkg/services/user"
	"github.com/grafana/grafana/pkg/tests/testsuite"
)

func TestMain(m *testing.M) {
	testsuite.Run(m)
}

func TestIntegrationAuthInfoStore(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping integration test")
	}

	sql := db.InitTestDB(t)
	store := ProvideStore(sql, secretstest.NewFakeSecretsService())

	t.Run("should be able to auth lables for users", func(t *testing.T) {
		ctx := context.Background()
		require.NoError(t, store.SetAuthInfo(ctx, &login.SetAuthInfoCommand{
			AuthModule: login.LDAPAuthModule,
			AuthId:     "1",
			UserId:     1,
		}))
		require.NoError(t, store.SetAuthInfo(ctx, &login.SetAuthInfoCommand{
			AuthModule: login.AzureADAuthModule,
			AuthId:     "1",
			UserId:     1,
		}))
		require.NoError(t, store.SetAuthInfo(ctx, &login.SetAuthInfoCommand{
			AuthModule: login.GoogleAuthModule,
			AuthId:     "10",
			UserId:     2,
		}))

		labels, err := store.GetUserLabels(ctx, login.GetUserLabelsQuery{UserIDs: []int64{1, 2}})
		require.NoError(t, err)
		require.Len(t, labels, 2)

		require.Equal(t, login.AzureADAuthModule, labels[1])
		require.Equal(t, login.GoogleAuthModule, labels[2])
	})

	t.Run("should always get the latest used", func(t *testing.T) {
		ctx := context.Background()
		require.NoError(t, store.SetAuthInfo(ctx, &login.SetAuthInfoCommand{
			AuthModule: login.LDAPAuthModule,
			AuthId:     "1",
			UserId:     1,
		}))

		defer func() {
			GetTime = time.Now
		}()

		GetTime = func() time.Time {
			return time.Now().Add(1 * time.Hour)
		}

		require.NoError(t, store.SetAuthInfo(ctx, &login.SetAuthInfoCommand{
			AuthModule: login.AzureADAuthModule,
			AuthId:     "2",
			UserId:     1,
		}))

		info, err := store.GetAuthInfo(ctx, &login.GetAuthInfoQuery{
			UserId: 1,
		})
		require.NoError(t, err)

		assert.Equal(t, login.AzureADAuthModule, info.AuthModule)
		assert.Equal(t, "2", info.AuthId)
	})

	t.Run("should return error when userID and authID is zero value", func(t *testing.T) {
		ctx := context.Background()
		info, err := store.GetAuthInfo(ctx, &login.GetAuthInfoQuery{
			AuthModule: login.GoogleAuthModule,
		})

		require.ErrorIs(t, err, user.ErrUserNotFound)
		require.Nil(t, info)
	})

	t.Run("should remove duplicates on update", func(t *testing.T) {
		ctx := context.Background()
		setCmd := &login.SetAuthInfoCommand{
			AuthModule: login.GenericOAuthModule,
			AuthId:     "1",
			UserId:     10,
		}

		require.NoError(t, store.SetAuthInfo(ctx, setCmd))
		require.NoError(t, store.SetAuthInfo(ctx, setCmd))

		count := countEntries(t, sql, setCmd.AuthModule, setCmd.AuthId, setCmd.UserId)
		require.Equal(t, 2, count)

		err := store.UpdateAuthInfo(ctx, &login.UpdateAuthInfoCommand{
			AuthModule: setCmd.AuthModule,
			AuthId:     setCmd.AuthId,
			UserId:     setCmd.UserId,
			OAuthToken: &oauth2.Token{
				AccessToken:  "atoken",
				RefreshToken: "rtoken",
				Expiry:       time.Now(),
			},
		})
		require.NoError(t, err)

		count = countEntries(t, sql, setCmd.AuthModule, setCmd.AuthId, setCmd.UserId)
		require.Equal(t, 1, count)
	})
}

func countEntries(t *testing.T, sql db.DB, authModule, authID string, userID int64) int {
	var result int

	err := sql.WithDbSession(context.Background(), func(sess *db.Session) error {
		_, err := sess.SQL(
			"SELECT COUNT(*) FROM user_auth WHERE auth_module = ? AND auth_id = ? AND user_id = ?",
			authModule, authID, userID,
		).Get(&result)
		return err
	})

	require.NoError(t, err)
	return result
}