xinmi/internal/application/service/digital_human_wechat_test.go

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)
}