package service import ( "context" "encoding/json" "testing" "time" "github.com/Tencent/WeKnora/internal/accesspolicy" "github.com/Tencent/WeKnora/internal/application/repository" "github.com/Tencent/WeKnora/internal/im" "github.com/Tencent/WeKnora/internal/im/wechat" "github.com/Tencent/WeKnora/internal/types" "github.com/google/uuid" "github.com/stretchr/testify/require" "gorm.io/driver/sqlite" "gorm.io/gorm" ) type digitalHumanTestChannels struct { items map[string]im.IMChannel created int } func (c *digitalHumanTestChannels) InvalidateDeletedChannels([]string) {} func (c *digitalHumanTestChannels) ListChannelsByAgent(agent string, tenant uint64) ([]im.IMChannel, error) { result := []im.IMChannel{} for _, item := range c.items { if item.AgentID == agent && item.TenantID == tenant { result = append(result, item) } } return result, nil } func (c *digitalHumanTestChannels) GetChannelByIDAndTenant(id string, tenant uint64) (*im.IMChannel, error) { item, exists := c.items[id] if !exists || item.TenantID != tenant { return nil, gorm.ErrRecordNotFound } return &item, nil } func (c *digitalHumanTestChannels) CreateChannel(item *im.IMChannel) error { item.ID = uuid.NewString() c.items[item.ID] = *item c.created++ return nil } func (c *digitalHumanTestChannels) UpdateChannel(item *im.IMChannel) error { c.items[item.ID] = *item return nil } func (c *digitalHumanTestChannels) DeleteChannel(id string, _ uint64) error { delete(c.items, id) return nil } type digitalHumanTestQR struct{ calls int } func (q *digitalHumanTestQR) GetLoginQRCode(context.Context) (*wechat.QRCodeResult, error) { return &wechat.QRCodeResult{QRCode: "upstream-ticket", QRCodeURL: "https://example.com/login"}, nil } func (q *digitalHumanTestQR) PollQRCodeStatus(context.Context, string) (*wechat.LoginResult, error) { q.calls++ return &wechat.LoginResult{Status: "confirmed", BotToken: "private-bot-token", ILinkBotID: "bot-id", ILinkUserID: "wechat-user"}, nil } func setupDigitalHumanWeChat(t *testing.T) (*DigitalHumanService, context.Context, *types.DigitalHuman, *digitalHumanTestChannels, *digitalHumanTestQR) { t.Helper() 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.AutoMigrate(&types.CustomAgent{}, &types.DigitalHuman{}, &types.DigitalHumanVersion{}, &types.DigitalHumanBinding{}, &types.DigitalHumanJob{})) require.NoError(t, db.Exec("CREATE TABLE users (id TEXT PRIMARY KEY, username TEXT, deleted_at DATETIME)").Error) require.NoError(t, db.Exec("CREATE TABLE im_channels (id TEXT PRIMARY KEY, agent_id TEXT, tenant_id INTEGER, enabled BOOLEAN, updated_at DATETIME, deleted_at DATETIME)").Error) repo := repository.NewDigitalHumanRepository(db) human := &types.DigitalHuman{ID: uuid.NewString(), TenantID: 7, OwnerID: "owner", Name: "Teacher", PrimaryAgentID: uuid.NewString(), LatestVersion: 1, Revision: 1, Status: "draft"} require.NoError(t, repo.Create(context.Background(), human, map[string]string{"tone": "clear"})) ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(7)) ctx = context.WithValue(ctx, types.UserIDContextKey, "owner") channels := &digitalHumanTestChannels{items: map[string]im.IMChannel{}} qr := &digitalHumanTestQR{} return &DigitalHumanService{repo: repo, channels: channels, qr: qr}, ctx, human, channels, qr } func TestDigitalHumanWeChatScopeAndIdempotency(t *testing.T) { service, ctx, human, channels, qr := setupDigitalHumanWeChat(t) ticket, err := service.WeChatQRCode(ctx, human.ID) require.NoError(t, err) require.NotEqual(t, "upstream-ticket", ticket.QRCode) for _, other := range []context.Context{ context.WithValue(ctx, types.UserIDContextKey, "other"), context.WithValue(ctx, types.TenantIDContextKey, uint64(8)), context.WithValue(context.WithValue(ctx, types.UserIDContextKey, "admin"), types.TenantRoleContextKey, types.TenantRoleAdmin), accesspolicy.WithSnapshot(context.WithValue(ctx, types.UserIDContextKey, "shared-reader"), &accesspolicy.Snapshot{Grants: map[string]bool{"digital-humans:" + human.ID: true}}), } { require.ErrorIs(t, service.CheckOwner(other, human.ID), types.ErrDigitalHumanNotFound) _, err := service.WeChatQRCode(other, human.ID) require.ErrorIs(t, err, types.ErrDigitalHumanNotFound) _, err = service.WeChatPoll(other, human.ID, ticket.QRCode) require.ErrorIs(t, err, types.ErrDigitalHumanNotFound) _, err = service.WeChatChannels(other, human.ID) require.ErrorIs(t, err, types.ErrDigitalHumanNotFound) } // A valid owner cannot use a ticket issued for another human. service.qrLogins[ticket.QRCode].humanID = "other-human" _, err = service.WeChatPoll(ctx, human.ID, ticket.QRCode) require.Error(t, err) service.qrLogins[ticket.QRCode].humanID = human.ID result, err := service.WeChatPoll(ctx, human.ID, ticket.QRCode) require.NoError(t, err) require.Equal(t, human.PrimaryAgentID, result.Channel.AgentID) _, err = service.WeChatPoll(ctx, human.ID, ticket.QRCode) require.NoError(t, err) require.Equal(t, 1, channels.created) require.Equal(t, 1, qr.calls) encoded, err := json.Marshal(result) require.NoError(t, err) require.NotContains(t, string(encoded), "private-bot-token") require.NotContains(t, string(encoded), "credentials\"") list, err := service.WeChatChannels(ctx, human.ID) require.NoError(t, err) require.Len(t, list, 1) sharedReader := accesspolicy.WithSnapshot(context.WithValue(ctx, types.UserIDContextKey, "shared-reader"), &accesspolicy.Snapshot{Grants: map[string]bool{"digital-humans:" + human.ID: true}}) require.ErrorIs(t, service.SetWeChatChannelEnabled(sharedReader, human.ID, result.Channel.ID, false), types.ErrDigitalHumanNotFound) require.ErrorIs(t, service.DeleteWeChatChannel(sharedReader, human.ID, result.Channel.ID), types.ErrDigitalHumanNotFound) require.True(t, channels.items[result.Channel.ID].Enabled) require.NoError(t, service.SetWeChatChannelEnabled(ctx, human.ID, result.Channel.ID, false)) require.False(t, channels.items[result.Channel.ID].Enabled) require.NoError(t, service.DeleteWeChatChannel(ctx, human.ID, result.Channel.ID)) require.Empty(t, channels.items) _, err = service.WeChatPoll(ctx, human.ID, ticket.QRCode) require.ErrorIs(t, err, types.ErrDigitalHumanNotFound) require.Equal(t, 1, channels.created) } func TestDigitalHumanWeChatRejectsExpiredAndForeignChannels(t *testing.T) { service, ctx, human, channels, qr := setupDigitalHumanWeChat(t) ticket, err := service.WeChatQRCode(ctx, human.ID) require.NoError(t, err) service.qrLogins[ticket.QRCode].expires = time.Now().Add(-time.Minute) _, err = service.WeChatPoll(ctx, human.ID, ticket.QRCode) require.Error(t, err) require.Zero(t, qr.calls) for _, channel := range []im.IMChannel{ {ID: "foreign-agent", TenantID: 7, AgentID: "other", Platform: "wechat"}, {ID: "foreign-tenant", TenantID: 8, AgentID: human.PrimaryAgentID, Platform: "wechat"}, {ID: "foreign-platform", TenantID: 7, AgentID: human.PrimaryAgentID, Platform: "feishu"}, } { channels.items[channel.ID] = channel require.ErrorIs(t, service.DeleteWeChatChannel(ctx, human.ID, channel.ID), types.ErrDigitalHumanNotFound) require.ErrorIs(t, service.SetWeChatChannelEnabled(ctx, human.ID, channel.ID, false), types.ErrDigitalHumanNotFound) } require.Len(t, channels.items, 3) }