grafana/pkg/infra/usagestats/statscollector/concurrent_users_test.go

package statscollector

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/grafana/grafana/pkg/infra/db"
	"github.com/grafana/grafana/pkg/services/stats/statsimpl"
	"github.com/grafana/grafana/pkg/setting"
	"github.com/grafana/grafana/pkg/tests/testsuite"
	"github.com/grafana/grafana/pkg/util"
)

// run tests with cleanup
func TestMain(m *testing.M) {
	testsuite.Run(m)
}

func TestConcurrentUsersMetrics(t *testing.T) {
	sqlStore, cfg := db.InitTestDBWithCfg(t)
	statsService := statsimpl.ProvideService(&setting.Cfg{}, sqlStore)
	s := createService(t, cfg, sqlStore, statsService)

	createConcurrentTokens(t, sqlStore)

	stats, err := s.collectConcurrentUsers(context.Background())
	require.NoError(t, err)

	assert.Equal(t, int32(1), stats["stats.auth_token_per_user_le_3"])
	assert.Equal(t, int32(2), stats["stats.auth_token_per_user_le_6"])
	assert.Equal(t, int32(3), stats["stats.auth_token_per_user_le_9"])
	assert.Equal(t, int32(4), stats["stats.auth_token_per_user_le_12"])
	assert.Equal(t, int32(5), stats["stats.auth_token_per_user_le_15"])
	assert.Equal(t, int32(6), stats["stats.auth_token_per_user_le_inf"])
}

func TestConcurrentUsersStats(t *testing.T) {
	sqlStore, cfg := db.InitTestDBWithCfg(t)
	statsService := statsimpl.ProvideService(&setting.Cfg{}, sqlStore)
	s := createService(t, cfg, sqlStore, statsService)

	createConcurrentTokens(t, sqlStore)

	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	t.Cleanup(func() {
		cancel()
	})

	actualResult, err := s.concurrentUsers(ctx)
	require.NoError(t, err)

	expectedResult := &concurrentUsersStats{
		BucketLE3:   1,
		BucketLE6:   2,
		BucketLE9:   3,
		BucketLE12:  4,
		BucketLE15:  5,
		BucketLEInf: 6,
	}
	assert.Equal(t, expectedResult, actualResult)

	createToken(t, 8, sqlStore)
	require.NoError(t, err)

	// Cached, new token is not reflected yet.
	actualResult, err = s.concurrentUsers(ctx)
	require.NoError(t, err)
	assert.Equal(t, expectedResult, actualResult)

	// Bust cache
	s.concurrentUserStatsCache = memoConcurrentUserStats{}

	expectedResult = &concurrentUsersStats{
		BucketLE3:   2,
		BucketLE6:   3,
		BucketLE9:   4,
		BucketLE12:  5,
		BucketLE15:  6,
		BucketLEInf: 7,
	}
	actualResult, err = s.concurrentUsers(ctx)
	require.NoError(t, err)
	assert.Equal(t, expectedResult, actualResult)
}

func createConcurrentTokens(t *testing.T, sqlStore db.DB) {
	t.Helper()
	for u := 1; u <= 6; u++ {
		for tkn := 1; tkn <= u*3; tkn++ {
			createToken(t, u, sqlStore)
		}
	}
}

func createToken(t *testing.T, uID int, sqlStore db.DB) {
	t.Helper()
	token, err := util.RandomHex(16)
	require.NoError(t, err)

	tokenWithSecret := fmt.Sprintf("%ssecret%d", token, uID)
	hashBytes := sha256.Sum256([]byte(tokenWithSecret))
	hashedToken := hex.EncodeToString(hashBytes[:])

	now := time.Now().Unix()

	userAuthToken := userAuthToken{
		UserID:        int64(uID),
		AuthToken:     hashedToken,
		PrevAuthToken: hashedToken,
		ClientIP:      "192.168.10.11",
		UserAgent:     "Mozilla",
		RotatedAt:     now,
		CreatedAt:     now,
		UpdatedAt:     now,
		SeenAt:        0,
		AuthTokenSeen: false,
	}

	err = sqlStore.WithDbSession(context.Background(), func(dbSession *db.Session) error {
		_, err = dbSession.Insert(&userAuthToken)
		return err
	})
	require.NoError(t, err)
}

type userAuthToken struct {
	UserID        int64 `xorm:"user_id"`
	AuthToken     string
	PrevAuthToken string
	UserAgent     string
	ClientIP      string `xorm:"client_ip"`
	AuthTokenSeen bool
	SeenAt        int64
	RotatedAt     int64
	CreatedAt     int64
	UpdatedAt     int64
}