heicode-mananger/heicode/controller/swarm_model_key_test.go

package controller

import (
	"fmt"
	"strings"
	"testing"

	"github.com/glebarez/sqlite"
	"github.com/stretchr/testify/require"
	"gorm.io/gorm"

	"github.com/heicode/manager/common"
	"github.com/heicode/manager/model"
)

func setupSwarmModelKeyTestDB(t *testing.T) {
	t.Helper()
	common.UsingSQLite = true
	common.UsingMySQL = false
	common.UsingPostgreSQL = false
	common.RedisEnabled = false

	dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
	db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
	require.NoError(t, err)
	model.DB = db
	require.NoError(t, db.AutoMigrate(&model.Token{}))
	t.Cleanup(func() {
		if sqlDB, err := db.DB(); err == nil {
			_ = sqlDB.Close()
		}
		model.DB = nil
	})
}

func TestSwarmModelKeyNames(t *testing.T) {
	require.Equal(t, "swarm:user:42", swarmModelKeyTokenName(42))
	require.Equal(t, "swarm-model-key-42", swarmModelKeySecretName(42))
}

func TestGetOrMintSwarmModelTokenReusesPerUser(t *testing.T) {
	setupSwarmModelKeyTestDB(t)

	// A.1: first call mints, second call reuses the SAME sk- across runs.
	first, err := getOrMintSwarmModelToken(7001)
	require.NoError(t, err)
	require.True(t, strings.HasPrefix(first, "sk-"))
	second, err := getOrMintSwarmModelToken(7001)
	require.NoError(t, err)
	require.Equal(t, first, second, "per-user swarm key must be reused across runs")

	// only one token persisted for the user, and it is hidden + system-managed.
	var toks []model.Token
	require.NoError(t, model.DB.Where("user_id = ?", 7001).Find(&toks).Error)
	require.Len(t, toks, 1)
	require.Equal(t, swarmModelKeyTokenName(7001), toks[0].Name)
	require.True(t, toks[0].HideFromUserUI)
	require.True(t, toks[0].UnlimitedQuota)
	require.EqualValues(t, -1, toks[0].ExpiredTime)

	// distinct users get distinct keys.
	other, err := getOrMintSwarmModelToken(7002)
	require.NoError(t, err)
	require.NotEqual(t, first, other)

	// invalid user rejected.
	_, err = getOrMintSwarmModelToken(0)
	require.Error(t, err)
}

func TestRevokeSwarmModelKeyReMintsFresh(t *testing.T) {
	setupSwarmModelKeyTestDB(t)

	first, err := getOrMintSwarmModelToken(7003)
	require.NoError(t, err)

	// revoke soft-deletes the token (KV not configured -> delete logged, non-fatal).
	revokeSwarmModelKey(7003)
	var live []model.Token
	require.NoError(t, model.DB.Where("user_id = ?", 7003).Find(&live).Error)
	require.Len(t, live, 0, "revoked swarm token must be soft-deleted")

	// next provision mints a brand-new key (A.5: stop -> revoke -> new key on re-run).
	second, err := getOrMintSwarmModelToken(7003)
	require.NoError(t, err)
	require.NotEqual(t, first, second)
}

func TestHandleSwarmPoolTerminatedRevokesByPayloadUserID(t *testing.T) {
	setupSwarmModelKeyTestDB(t)
	_, err := getOrMintSwarmModelToken(7004)
	require.NoError(t, err)

	payload := agentCallbackEnvelope{
		EventType: "swarm.pool_terminated",
		Payload:   map[string]any{"user_id": "7004", "secret_ref": "azkv://heicode-kv.vault.azure.net/secrets/swarm-model-key-7004"},
	}
	handleSwarmPoolTerminated(payload, agentDeploymentRecord{})

	var live []model.Token
	require.NoError(t, model.DB.Where("user_id = ?", 7004).Find(&live).Error)
	require.Len(t, live, 0)
}

func TestHandleSwarmPoolTerminatedFallsBackToDeploymentUser(t *testing.T) {
	setupSwarmModelKeyTestDB(t)
	_, err := getOrMintSwarmModelToken(7005)
	require.NoError(t, err)

	// no user_id in payload -> fall back to deployment user context.
	payload := agentCallbackEnvelope{EventType: "swarm.pool_terminated", Payload: map[string]any{}}
	record := agentDeploymentRecord{}
	record.Plan.UserContext.UserID = "7005"
	handleSwarmPoolTerminated(payload, record)

	var live []model.Token
	require.NoError(t, model.DB.Where("user_id = ?", 7005).Find(&live).Error)
	require.Len(t, live, 0)
}

func TestHandleSwarmPoolTerminatedIgnoresOtherEvents(t *testing.T) {
	setupSwarmModelKeyTestDB(t)
	_, err := getOrMintSwarmModelToken(7006)
	require.NoError(t, err)

	// wrong event type -> no revocation.
	handleSwarmPoolTerminated(agentCallbackEnvelope{EventType: "swarm.stopped", Payload: map[string]any{"user_id": "7006"}}, agentDeploymentRecord{})
	var live []model.Token
	require.NoError(t, model.DB.Where("user_id = ?", 7006).Find(&live).Error)
	require.Len(t, live, 1)
}

func TestSwarmPoolTerminatedEventRegistered(t *testing.T) {
	// control-plane lifecycle event must be a known callback type so the schema
	// validator accepts it and routes it to the swarm_lifecycle category.
	_, hasFields := agentCallbackEventRequiredFields["swarm.pool_terminated"]
	require.True(t, hasFields)
	require.Equal(t, "swarm_lifecycle", agentCallbackEventCategories["swarm.pool_terminated"])
}