xinmi/internal/im/session_channel_isolation_test.go

package im

import (
	"context"
	"testing"

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

type channelIsolationSessionService struct {
	interfaces.SessionService
	db *gorm.DB
}

func (s *channelIsolationSessionService) CreateSession(ctx context.Context, session *types.Session) (*types.Session, error) {
	return session, s.db.WithContext(ctx).Create(session).Error
}

func TestSameAgentConversationsStayIsolatedBetweenConnectedAccounts(t *testing.T) {
	for _, mode := range []string{"user", "thread"} {
		t.Run(mode, func(t *testing.T) {
			db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
			require.NoError(t, err)
			require.NoError(t, db.AutoMigrate(&types.Session{}))
			require.NoError(t, db.Exec(`CREATE TABLE im_channel_sessions (
                id TEXT PRIMARY KEY, platform TEXT, user_id TEXT, chat_id TEXT,
                thread_id TEXT, session_id TEXT, tenant_id INTEGER, agent_id TEXT,
                im_channel_id TEXT, status TEXT, metadata TEXT, created_at DATETIME,
                updated_at DATETIME, deleted_at DATETIME
            )`).Error)
			svc := &Service{db: db, sessionService: &channelIsolationSessionService{db: db}}
			msg := &IncomingMessage{Platform: PlatformFeishu, UserID: "same-user", ChatID: "same-chat", ThreadID: "same-thread"}
			first, err := svc.resolveSession(context.Background(), msg, 1, "shared-agent", "account-a", mode)
			require.NoError(t, err)
			second, err := svc.resolveSession(context.Background(), msg, 1, "shared-agent", "account-b", mode)
			require.NoError(t, err)
			require.NotEqual(t, first.SessionID, second.SessionID)
			again, err := svc.resolveSession(context.Background(), msg, 1, "shared-agent", "account-a", mode)
			require.NoError(t, err)
			require.Equal(t, first.SessionID, again.SessionID)
		})
	}
}