xinmi/internal/im/personal_channel_access_test.go

package im

import (
	"context"
	"testing"

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

func TestDigitalHumanPersonalChannelAccessAcrossGenericIMEndpoints(t *testing.T) {
	db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	t.Cleanup(func() { sqlDB.Close() })
	require.NoError(t, db.Exec("CREATE TABLE digital_humans (id TEXT, primary_agent_id TEXT, tenant_id INTEGER, owner_id TEXT, deleted_at DATETIME)").Error)
	require.NoError(t, db.Exec("CREATE TABLE im_channels (id TEXT, agent_id TEXT, tenant_id INTEGER)").Error)
	require.NoError(t, db.Exec("INSERT INTO digital_humans (id, primary_agent_id, tenant_id, owner_id) VALUES ('human', 'personal', 7, 'owner')").Error)
	require.NoError(t, db.Exec("INSERT INTO im_channels VALUES ('personal-channel', 'personal', 7), ('shared-channel', 'shared', 7), ('different-tenant', 'personal', 8)").Error)
	service := &Service{db: db}
	owner := context.WithValue(context.Background(), types.UserIDContextKey, "owner")
	require.NoError(t, service.CheckPersonalChannelAgentAccess(owner, "personal", 7))
	for _, user := range []string{"other", "admin", "system-7", ""} {
		ctx := context.WithValue(context.Background(), types.UserIDContextKey, user)
		ctx = context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleAdmin)
		require.ErrorIs(t, service.CheckPersonalChannelAgentAccess(ctx, "personal", 7), gorm.ErrRecordNotFound)
		require.NoError(t, service.CheckPersonalChannelAgentAccess(ctx, "shared", 7))
		var ids []string
		require.NoError(t, excludeOtherPersonalChannels(db.Table("im_channels AS c").Where("c.tenant_id = 7"), ctx).Order("id").Pluck("id", &ids).Error)
		require.Equal(t, []string{"shared-channel"}, ids)
	}
	var ids []string
	require.NoError(t, excludeOtherPersonalChannels(db.Table("im_channels AS c").Where("c.tenant_id = 7"), owner).Order("id").Pluck("id", &ids).Error)
	require.Equal(t, []string{"personal-channel", "shared-channel"}, ids)
}