483 lines
15 KiB
Go
483 lines
15 KiB
Go
package ws
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/alicebob/miniredis/v2"
|
||
"github.com/redis/go-redis/v9"
|
||
"github.com/stretchr/testify/assert"
|
||
"github.com/stretchr/testify/require"
|
||
)
|
||
|
||
// --- 测试辅助函数 ---
|
||
// setupPresenceTest 创建 miniredis 实例、redis.Client、BroadcastRelay 和 PresenceTracker，
|
||
// 测试结束后需要调用 cleanup。
|
||
|
||
func setupPresenceTest(t *testing.T) (*miniredis.Miniredis, *redis.Client, *PresenceTracker, *mockHandler, func()) {
|
||
t.Helper()
|
||
mr := miniredis.RunT(t)
|
||
rdb := redis.NewClient(&redis.Options{
|
||
Addr: mr.Addr(),
|
||
})
|
||
handler := newMockHandler()
|
||
relay := NewBroadcastRelay(rdb, handler)
|
||
tracker := NewPresenceTracker(rdb, relay)
|
||
cleanup := func() {
|
||
rdb.Close()
|
||
mr.Close()
|
||
}
|
||
return mr, rdb, tracker, handler, cleanup
|
||
}
|
||
|
||
// --- PresenceTracker 构造函数测试 ---
|
||
|
||
func TestNewPresenceTracker(t *testing.T) {
|
||
_, rdb, _, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
handler := newMockHandler()
|
||
relay := NewBroadcastRelay(rdb, handler)
|
||
tracker := NewPresenceTracker(rdb, relay)
|
||
|
||
assert.NotNil(t, tracker, "PresenceTracker 应成功创建")
|
||
assert.Equal(t, rdb, tracker.rdb, "Redis 客户端应正确设置")
|
||
assert.Equal(t, relay, tracker.relay, "BroadcastRelay 应正确设置")
|
||
}
|
||
|
||
// --- SetAgentOnline 测试 ---
|
||
|
||
func TestSetAgentOnline(t *testing.T) {
|
||
_, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
agentID := uint(42)
|
||
accountID := uint(1)
|
||
|
||
err := tracker.SetAgentOnline(ctx, agentID, accountID)
|
||
require.NoError(t, err, "SetAgentOnline 应成功执行")
|
||
|
||
// 验证 Redis sorted set 中存在该成员
|
||
member := fmt.Sprintf("%d:%d", agentID, accountID)
|
||
score, err := rdb.ZScore(ctx, RedisKeyPresenceAgents, member).Result()
|
||
require.NoError(t, err, "应能从 sorted set 中获取成员 score")
|
||
// score 应是当前 Unix 时间戳
|
||
now := float64(time.Now().Unix())
|
||
assert.InDelta(t, now, score, 2.0, "score 应接近当前时间戳")
|
||
|
||
// 验证 hash 中存在该成员的状态
|
||
status, err := rdb.HGet(ctx, RedisKeyPresenceStatus, member).Result()
|
||
require.NoError(t, err, "应能从 hash 中获取成员状态")
|
||
assert.Equal(t, "online", status, "状态应为 online")
|
||
}
|
||
|
||
func TestSetAgentOnline_MultipleAgents(t *testing.T) {
|
||
_, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 设置多个 agent 在线
|
||
err := tracker.SetAgentOnline(ctx, 1, 10)
|
||
require.NoError(t, err)
|
||
err = tracker.SetAgentOnline(ctx, 2, 10)
|
||
require.NoError(t, err)
|
||
err = tracker.SetAgentOnline(ctx, 3, 20)
|
||
require.NoError(t, err)
|
||
|
||
// 验证 sorted set 中有3个成员
|
||
count, err := rdb.ZCard(ctx, RedisKeyPresenceAgents).Result()
|
||
require.NoError(t, err)
|
||
assert.Equal(t, int64(3), count, "应有3个 agent 在线")
|
||
}
|
||
|
||
func TestSetAgentOnline_RefreshExisting(t *testing.T) {
|
||
mr, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 先设置在线
|
||
err := tracker.SetAgentOnline(ctx, 5, 1)
|
||
require.NoError(t, err)
|
||
|
||
// 等一小段时间再刷新
|
||
mr.FastForward(2 * time.Second)
|
||
|
||
// 再次设置在线（模拟心跳刷新）
|
||
err = tracker.SetAgentOnline(ctx, 5, 1)
|
||
require.NoError(t, err, "刷新已有 agent 应成功")
|
||
|
||
// 验证 score 已更新
|
||
member := "5:1"
|
||
score, err := rdb.ZScore(ctx, RedisKeyPresenceAgents, member).Result()
|
||
require.NoError(t, err)
|
||
now := float64(time.Now().Unix())
|
||
assert.InDelta(t, now, score, 2.0, "刷新后 score 应更新为当前时间戳")
|
||
}
|
||
|
||
// --- SetAgentOffline 测试 ---
|
||
|
||
func TestSetAgentOffline(t *testing.T) {
|
||
_, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
agentID := uint(42)
|
||
accountID := uint(1)
|
||
|
||
// 先设置在线
|
||
err := tracker.SetAgentOnline(ctx, agentID, accountID)
|
||
require.NoError(t, err)
|
||
|
||
// 然后设置离线
|
||
err = tracker.SetAgentOffline(ctx, agentID, accountID)
|
||
require.NoError(t, err, "SetAgentOffline 应成功执行")
|
||
|
||
// 验证 sorted set 中不再有该成员
|
||
member := fmt.Sprintf("%d:%d", agentID, accountID)
|
||
_, err = rdb.ZScore(ctx, RedisKeyPresenceAgents, member).Result()
|
||
assert.Error(t, err, "离线后 sorted set 中不应有该成员")
|
||
assert.Equal(t, redis.Nil, err, "应返回 redis.Nil")
|
||
|
||
// 验证 hash 中不再有该成员的状态
|
||
_, err = rdb.HGet(ctx, RedisKeyPresenceStatus, member).Result()
|
||
assert.Equal(t, redis.Nil, err, "hash 中不应有该成员的状态")
|
||
}
|
||
|
||
func TestSetAgentOffline_NotOnline(t *testing.T) {
|
||
_, _, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 设置从未在线的 agent 为离线 — 应不会出错
|
||
err := tracker.SetAgentOffline(ctx, 99, 1)
|
||
require.NoError(t, err, "对不存在的 agent 设置离线不应报错")
|
||
}
|
||
|
||
// --- SetAgentBusy 测试 ---
|
||
|
||
func TestSetAgentBusy(t *testing.T) {
|
||
_, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 先设置在线
|
||
err := tracker.SetAgentOnline(ctx, 10, 5)
|
||
require.NoError(t, err)
|
||
|
||
// 设置为 busy
|
||
err = tracker.SetAgentBusy(ctx, 10, 5)
|
||
require.NoError(t, err, "SetAgentBusy 应成功执行")
|
||
|
||
// 验证 hash 中状态为 busy
|
||
member := "10:5"
|
||
status, err := rdb.HGet(ctx, RedisKeyPresenceStatus, member).Result()
|
||
require.NoError(t, err)
|
||
assert.Equal(t, "busy", status, "状态应为 busy")
|
||
|
||
// 验证仍在 sorted set 中（busy 意味着仍在线）
|
||
score, err := rdb.ZScore(ctx, RedisKeyPresenceAgents, member).Result()
|
||
require.NoError(t, err)
|
||
assert.Greater(t, score, 0.0, "busy agent 应仍在 sorted set 中")
|
||
}
|
||
|
||
// --- SetContactOnline / SetContactOffline 测试 ---
|
||
|
||
func TestSetContactOnline(t *testing.T) {
|
||
_, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
contactID := uint(100)
|
||
accountID := uint(2)
|
||
|
||
err := tracker.SetContactOnline(ctx, contactID, accountID)
|
||
require.NoError(t, err, "SetContactOnline 应成功执行")
|
||
|
||
member := fmt.Sprintf("%d:%d", contactID, accountID)
|
||
score, err := rdb.ZScore(ctx, RedisKeyPresenceContacts, member).Result()
|
||
require.NoError(t, err)
|
||
now := float64(time.Now().Unix())
|
||
assert.InDelta(t, now, score, 2.0, "score 应接近当前时间戳")
|
||
|
||
status, err := rdb.HGet(ctx, RedisKeyPresenceStatus, member).Result()
|
||
require.NoError(t, err)
|
||
assert.Equal(t, "online", status, "状态应为 online")
|
||
}
|
||
|
||
func TestSetContactOffline(t *testing.T) {
|
||
_, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 先设置在线
|
||
err := tracker.SetContactOnline(ctx, 100, 2)
|
||
require.NoError(t, err)
|
||
|
||
// 设置离线
|
||
err = tracker.SetContactOffline(ctx, 100, 2)
|
||
require.NoError(t, err, "SetContactOffline 应成功执行")
|
||
|
||
// 验证离线后成员已被移除
|
||
assert.Equal(t, redis.Nil, rdb.ZScore(ctx, RedisKeyPresenceContacts, "100:2").Err(), "离线后 sorted set 中不应有该成员")
|
||
assert.Equal(t, redis.Nil, rdb.HGet(ctx, RedisKeyPresenceStatus, "100:2").Err(), "hash 中不应有该成员的状态")
|
||
}
|
||
|
||
// --- GetOnlineAgentsForAccount 测试 ---
|
||
|
||
func TestGetOnlineAgentsForAccount(t *testing.T) {
|
||
_, _, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 设置 account=10 的3个 agent 在线
|
||
err := tracker.SetAgentOnline(ctx, 1, 10)
|
||
require.NoError(t, err)
|
||
err = tracker.SetAgentOnline(ctx, 2, 10)
|
||
require.NoError(t, err)
|
||
err = tracker.SetAgentOnline(ctx, 3, 10)
|
||
require.NoError(t, err)
|
||
|
||
// 设置 account=20 的1个 agent 在线
|
||
err = tracker.SetAgentOnline(ctx, 4, 20)
|
||
require.NoError(t, err)
|
||
|
||
// 获取 account=10 的在线 agent
|
||
agents, err := tracker.GetOnlineAgentsForAccount(ctx, 10)
|
||
require.NoError(t, err, "GetOnlineAgentsForAccount 应成功执行")
|
||
assert.Len(t, agents, 3, "account=10 应有3个在线 agent")
|
||
|
||
// 验证返回的 ID 包含 1, 2, 3
|
||
assert.Contains(t, agents, uint(1))
|
||
assert.Contains(t, agents, uint(2))
|
||
assert.Contains(t, agents, uint(3))
|
||
|
||
// 获取 account=20 的在线 agent
|
||
agents, err = tracker.GetOnlineAgentsForAccount(ctx, 20)
|
||
require.NoError(t, err)
|
||
assert.Len(t, agents, 1, "account=20 应有1个在线 agent")
|
||
assert.Contains(t, agents, uint(4))
|
||
|
||
// 获取无在线 agent 的 account
|
||
agents, err = tracker.GetOnlineAgentsForAccount(ctx, 30)
|
||
require.NoError(t, err)
|
||
assert.Empty(t, agents, "account=30 不应有在线 agent")
|
||
}
|
||
|
||
func TestGetOnlineAgentsForAccount_Expired(t *testing.T) {
|
||
_, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 设置 agent 在线
|
||
err := tracker.SetAgentOnline(ctx, 1, 10)
|
||
require.NoError(t, err)
|
||
|
||
// 直接修改 sorted set 的 score 为过去的时间戳，模拟过期
|
||
// SetAgentOnline 会设置 score 为当前时间戳，我们需要将它改为一个
|
||
// 低于阈值（time.Now().Unix() - PresenceDurationAgentSec）的值
|
||
member := "1:10"
|
||
expiredScore := float64(time.Now().Unix() - PresenceDurationAgentSec - 10)
|
||
err = rdb.ZAdd(ctx, RedisKeyPresenceAgents, redis.Z{Score: expiredScore, Member: member}).Err()
|
||
require.NoError(t, err, "应能修改 agent score 为过期值")
|
||
|
||
// 同时更新 hash 中的状态时间戳，使状态与 score 一致
|
||
err = rdb.HSet(ctx, RedisKeyPresenceStatus, member, "online").Err()
|
||
require.NoError(t, err)
|
||
|
||
// 获取在线 agent — 过期的 agent 不应被返回
|
||
agents, err := tracker.GetOnlineAgentsForAccount(ctx, 10)
|
||
require.NoError(t, err)
|
||
assert.Empty(t, agents, "过期 agent 不应在在线列表中")
|
||
}
|
||
|
||
// --- GetAgentStatus 测试 ---
|
||
|
||
func TestGetAgentStatus(t *testing.T) {
|
||
_, _, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 设置 online
|
||
err := tracker.SetAgentOnline(ctx, 5, 1)
|
||
require.NoError(t, err)
|
||
|
||
status, err := tracker.GetAgentStatus(ctx, 5, 1)
|
||
require.NoError(t, err)
|
||
assert.Equal(t, "online", status, "状态应为 online")
|
||
|
||
// 设置 busy
|
||
err = tracker.SetAgentBusy(ctx, 5, 1)
|
||
require.NoError(t, err)
|
||
|
||
status, err = tracker.GetAgentStatus(ctx, 5, 1)
|
||
require.NoError(t, err)
|
||
assert.Equal(t, "busy", status, "状态应为 busy")
|
||
|
||
// 设置 offline
|
||
err = tracker.SetAgentOffline(ctx, 5, 1)
|
||
require.NoError(t, err)
|
||
|
||
status, err = tracker.GetAgentStatus(ctx, 5, 1)
|
||
require.NoError(t, err)
|
||
assert.Equal(t, "offline", status, "离线 agent 状态应为 offline")
|
||
}
|
||
|
||
func TestGetAgentStatus_UnknownAgent(t *testing.T) {
|
||
_, _, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 未知 agent — 应返回 offline
|
||
status, err := tracker.GetAgentStatus(ctx, 999, 1)
|
||
require.NoError(t, err)
|
||
assert.Equal(t, "offline", status, "未知 agent 状态应为 offline")
|
||
}
|
||
|
||
// --- RefreshAgentPresence / RefreshContactPresence 测试 ---
|
||
|
||
func TestRefreshAgentPresence(t *testing.T) {
|
||
mr, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 先设置在线
|
||
err := tracker.SetAgentOnline(ctx, 5, 1)
|
||
require.NoError(t, err)
|
||
|
||
// 推进一小段时间
|
||
mr.FastForward(3 * time.Second)
|
||
|
||
// 刷新（模拟心跳）
|
||
err = tracker.RefreshAgentPresence(ctx, 5, 1)
|
||
require.NoError(t, err, "RefreshAgentPresence 应成功执行")
|
||
|
||
member := "5:1"
|
||
score, err := rdb.ZScore(ctx, RedisKeyPresenceAgents, member).Result()
|
||
require.NoError(t, err)
|
||
now := float64(time.Now().Unix())
|
||
assert.InDelta(t, now, score, 2.0, "刷新后 score 应更新")
|
||
}
|
||
|
||
func TestRefreshContactPresence(t *testing.T) {
|
||
mr, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 先设置在线
|
||
err := tracker.SetContactOnline(ctx, 100, 2)
|
||
require.NoError(t, err)
|
||
|
||
// 推进一小段时间
|
||
mr.FastForward(5 * time.Second)
|
||
|
||
// 刷新
|
||
err = tracker.RefreshContactPresence(ctx, 100, 2)
|
||
require.NoError(t, err, "RefreshContactPresence 应成功执行")
|
||
|
||
member := "100:2"
|
||
score, err := rdb.ZScore(ctx, RedisKeyPresenceContacts, member).Result()
|
||
require.NoError(t, err)
|
||
now := float64(time.Now().Unix())
|
||
assert.InDelta(t, now, score, 2.0, "刷新后 score 应更新")
|
||
}
|
||
|
||
// --- CleanupExpired 测试 ---
|
||
|
||
func TestCleanupExpired(t *testing.T) {
|
||
_, rdb, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 设置 agent 和 contact 在线（score 为当前时间戳）
|
||
err := tracker.SetAgentOnline(ctx, 1, 10)
|
||
require.NoError(t, err)
|
||
err = tracker.SetContactOnline(ctx, 100, 10)
|
||
require.NoError(t, err)
|
||
|
||
// 直接修改 sorted set score 为过期值，模拟 agent 已过期
|
||
// agent 过期阈值: score < (time.Now().Unix() - PresenceDurationAgentSec)
|
||
expiredAgentScore := float64(time.Now().Unix() - PresenceDurationAgentSec - 10)
|
||
err = rdb.ZAdd(ctx, RedisKeyPresenceAgents, redis.Z{Score: expiredAgentScore, Member: "1:10"}).Err()
|
||
require.NoError(t, err, "应能修改 agent score 为过期值")
|
||
|
||
// contact 的 score 仍然在有效范围内（当前时间戳），不会被清理
|
||
// 执行清理
|
||
err = tracker.CleanupExpired(ctx)
|
||
require.NoError(t, err, "CleanupExpired 应成功执行")
|
||
|
||
// 验证过期 agent 已被清除
|
||
agentCount, err := rdb.ZCard(ctx, RedisKeyPresenceAgents).Result()
|
||
require.NoError(t, err)
|
||
assert.Equal(t, int64(0), agentCount, "过期 agent 应被清理")
|
||
|
||
// contact 仍在有效期内，不应被清除
|
||
contactCount, err := rdb.ZCard(ctx, RedisKeyPresenceContacts).Result()
|
||
require.NoError(t, err)
|
||
assert.Equal(t, int64(1), contactCount, "有效期内 contact 不应被清理")
|
||
|
||
// 现在让 contact 也过期
|
||
expiredContactScore := float64(time.Now().Unix() - PresenceDurationContactSec - 10)
|
||
err = rdb.ZAdd(ctx, RedisKeyPresenceContacts, redis.Z{Score: expiredContactScore, Member: "100:10"}).Err()
|
||
require.NoError(t, err, "应能修改 contact score 为过期值")
|
||
|
||
// 再次清理
|
||
err = tracker.CleanupExpired(ctx)
|
||
require.NoError(t, err)
|
||
|
||
// 验证过期 contact 也被清除
|
||
contactCount, err = rdb.ZCard(ctx, RedisKeyPresenceContacts).Result()
|
||
require.NoError(t, err)
|
||
assert.Equal(t, int64(0), contactCount, "过期 contact 应被清理")
|
||
}
|
||
|
||
func TestCleanupExpired_NoExpired(t *testing.T) {
|
||
_, _, tracker, _, cleanup := setupPresenceTest(t)
|
||
defer cleanup()
|
||
|
||
ctx := context.Background()
|
||
|
||
// 设置在线
|
||
err := tracker.SetAgentOnline(ctx, 1, 10)
|
||
require.NoError(t, err)
|
||
|
||
// 不推进时间，直接清理
|
||
err = tracker.CleanupExpired(ctx)
|
||
require.NoError(t, err, "没有过期数据时清理不应报错")
|
||
}
|
||
|
||
// --- presenceMember / parsePresenceMember 辅助函数测试 ---
|
||
|
||
func TestPresenceMember(t *testing.T) {
|
||
result := presenceMember(42, 1)
|
||
assert.Equal(t, "42:1", result, "应生成正确的成员字符串")
|
||
}
|
||
|
||
func TestParsePresenceMember(t *testing.T) {
|
||
id, acct := parsePresenceMember("42:1")
|
||
assert.Equal(t, uint(42), id, "应正确解析 ID")
|
||
assert.Equal(t, uint(1), acct, "应正确解析 account ID")
|
||
|
||
// 无效格式
|
||
id, acct = parsePresenceMember("invalid")
|
||
assert.Equal(t, uint(0), id, "无效格式应返回0")
|
||
assert.Equal(t, uint(0), acct, "无效格式应返回0")
|
||
|
||
// 缺少分隔符
|
||
id, acct = parsePresenceMember("42")
|
||
assert.Equal(t, uint(0), id, "缺少分隔符应返回0")
|
||
} |