xinmi/internal/accesscontrol/permission_version_migration_test.go

package accesscontrol

import (
	"context"
	"fmt"
	"net/http"
	"net/http/httptest"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/Tencent/WeKnora/internal/types"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

func TestRolePermissionVersionMigrationPreservesLegacyAndAllowsGranularSave(t *testing.T) {
	for _, dialect := range []string{"sqlite", "postgres"} {
		t.Run(dialect, func(t *testing.T) {
			db := permissionMigrationDB(t, dialect)
			migrationDir := "sqlite"
			if dialect == "postgres" {
				migrationDir = "versioned"
				applyPermissionMigration(t, db, migrationDir, "000196_workbench_access.up.sql")
			} else {
				require.NoError(t, db.Exec(`CREATE TABLE workbench_access_roles (
id TEXT PRIMARY KEY, tenant_id INTEGER NOT NULL, name TEXT NOT NULL,
description TEXT NOT NULL DEFAULT '', permissions TEXT NOT NULL DEFAULT '[]',
version INTEGER NOT NULL DEFAULT 1, created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, UNIQUE (tenant_id, name))`).Error)
			}
			require.False(t, db.Migrator().HasColumn(&Role{}, "permission_version"))
			require.NoError(t, db.Exec(`INSERT INTO workbench_access_roles
(id, tenant_id, name, permissions, version) VALUES (?, ?, ?, ?, ?)`,
				"legacy", 1, "legacy", types.StringArray{"agents.read", "agents.write"}, 7).Error)

			applyPermissionMigration(t, db, migrationDir, "000215_access_role_permission_version.up.sql")
			var legacy Role
			require.NoError(t, db.First(&legacy, "id = ?", "legacy").Error)
			require.Equal(t, 0, legacy.PermissionVersion)
			require.Equal(t, 7, legacy.Version)
			require.Equal(t, types.StringArray{"agents.read", "agents.write"}, legacy.Permissions)
			require.Contains(t, effectiveRolePermissions(legacy), "agents.skills")
			require.Contains(t, effectiveRolePermissions(legacy), "agents.memory")

			// Only unrelated audit infrastructure uses AutoMigrate: the role schema
			// must come from the shipped SQL to catch production upgrade omissions.
			require.NoError(t, db.AutoMigrate(&types.AuditLog{}))
			s := NewService(db)
			router := gin.New()
			s.Register(router.Group("/api/v1"))
			for _, action := range []struct{ method, path, body string }{
				{"POST", "/api/v1/access/roles", `{"name":"new-granular","permissions":["agents.read","agents.write"]}`},
				{"PUT", "/api/v1/access/roles/legacy", `{"name":"legacy","permissions":["agents.read","agents.write"],"version":7}`},
			} {
				req := httptest.NewRequest(action.method, action.path, strings.NewReader(action.body)).WithContext(versionContext(1, "admin", true))
				req.Header.Set("Content-Type", "application/json")
				res := httptest.NewRecorder()
				router.ServeHTTP(res, req)
				require.Equal(t, http.StatusOK, res.Code, res.Body.String())
			}
			var roles []Role
			require.NoError(t, db.Order("name").Find(&roles).Error)
			require.Len(t, roles, 2)
			for _, role := range roles {
				require.Equal(t, currentPermissionVersion, role.PermissionVersion)
				require.NotContains(t, effectiveRolePermissions(role), "agents.skills")
				require.NotContains(t, effectiveRolePermissions(role), "agents.memory")
			}
			if dialect == "postgres" {
				// The production hotfix can precede the next deployment; rerunning
				// migration 215 must not reset already saved granular selections.
				applyPermissionMigration(t, db, migrationDir, "000215_access_role_permission_version.up.sql")
				require.NoError(t, db.First(&legacy, "id = ?", "legacy").Error)
				require.Equal(t, currentPermissionVersion, legacy.PermissionVersion)
				require.Equal(t, 8, legacy.Version)
			}
			applyPermissionMigration(t, db, migrationDir, "000215_access_role_permission_version.down.sql")
			require.False(t, db.Migrator().HasColumn(&Role{}, "permission_version"))
			var count int64
			require.NoError(t, db.Table("workbench_access_roles").Count(&count).Error)
			require.EqualValues(t, 2, count)
		})
	}
}

func applyPermissionMigration(t *testing.T, db *gorm.DB, dialect, name string) {
	t.Helper()
	sql, err := os.ReadFile(filepath.Join("..", "..", "migrations", dialect, name))
	require.NoError(t, err)
	require.NoError(t, db.Exec(string(sql)).Error)
}

func permissionMigrationDB(t *testing.T, dialect string) *gorm.DB {
	t.Helper()
	var dialector gorm.Dialector
	if dialect == "sqlite" {
		dialector = sqlite.Open(filepath.Join(t.TempDir(), "roles.db"))
	} else {
		bin := os.Getenv("POSTGRES_TEST_BIN")
		if bin == "" {
			bin = "/opt/homebrew/opt/postgresql@16/bin"
		}
		if _, err := os.Stat(filepath.Join(bin, "initdb")); err != nil {
			t.Skip("isolated PostgreSQL binaries unavailable; set POSTGRES_TEST_BIN")
		}
		dir, err := os.MkdirTemp("/tmp", "role-migration-pg-")
		require.NoError(t, err)
		run := func(name string, args ...string) {
			t.Helper()
			ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			defer cancel()
			output, err := exec.CommandContext(ctx, filepath.Join(bin, name), args...).CombinedOutput()
			require.NoError(t, err, "%s: %s", name, output)
		}
		data := filepath.Join(dir, "data")
		running := false
		t.Cleanup(func() {
			if running {
				run("pg_ctl", "-D", data, "-m", "fast", "-w", "stop")
			}
			require.NoError(t, os.RemoveAll(dir))
		})
		run("initdb", "-D", data, "-U", "fixture", "--auth=trust", "--no-locale", "--encoding=UTF8")
		run("pg_ctl", "-D", data, "-l", filepath.Join(dir, "postgres.log"), "-o",
			fmt.Sprintf("-F -c listen_addresses='' -k %s -p 55461", dir), "-w", "start")
		running = true
		dialector = postgres.Open(fmt.Sprintf("host=%s port=55461 user=fixture dbname=postgres sslmode=disable", dir))
	}
	db, err := gorm.Open(dialector, &gorm.Config{})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	t.Cleanup(func() { require.NoError(t, sqlDB.Close()) })
	return db
}