You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
327 lines
7.6 KiB
327 lines
7.6 KiB
package utils
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
// 测试用的 Redis 客户端
|
|
func setupTestRedis() *redis.Client {
|
|
return redis.NewClient(&redis.Options{
|
|
Addr: "localhost:6379",
|
|
Password: "",
|
|
DB: 1, // 使用测试数据库
|
|
})
|
|
}
|
|
|
|
// 测试 JWT 工具的基本功能
|
|
func TestJWTUtil_BasicOperations(t *testing.T) {
|
|
redisClient := setupTestRedis()
|
|
defer redisClient.Close()
|
|
|
|
// 清理测试数据
|
|
redisClient.FlushDB(context.Background())
|
|
|
|
jwtUtil := NewJWTUtil("test-secret", 3600, redisClient)
|
|
ctx := context.Background()
|
|
|
|
userID := int64(123)
|
|
username := "testuser"
|
|
|
|
// 测试生成 token
|
|
token, err := jwtUtil.GenerateToken(ctx, userID, username)
|
|
assert.NoError(t, err)
|
|
assert.NotEmpty(t, token)
|
|
|
|
// 测试验证 token
|
|
claims, err := jwtUtil.ValidateToken(ctx, token)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, userID, claims.UserID)
|
|
assert.Equal(t, username, claims.Username)
|
|
|
|
// 测试 token 是否存在
|
|
exists, err := jwtUtil.TokenExists(ctx, token)
|
|
assert.NoError(t, err)
|
|
assert.True(t, exists)
|
|
|
|
// 测试解析 token（不验证签名）
|
|
unverifiedClaims, err := jwtUtil.ParseTokenUnverified(token)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, userID, unverifiedClaims.UserID)
|
|
assert.Equal(t, username, unverifiedClaims.Username)
|
|
|
|
// 测试删除 token
|
|
err = jwtUtil.DeleteToken(ctx, token)
|
|
assert.NoError(t, err)
|
|
|
|
// 验证 token 已被删除
|
|
exists, err = jwtUtil.TokenExists(ctx, token)
|
|
assert.NoError(t, err)
|
|
assert.False(t, exists)
|
|
}
|
|
|
|
// 测试 token 刷新功能
|
|
func TestJWTUtil_RefreshToken(t *testing.T) {
|
|
redisClient := setupTestRedis()
|
|
defer redisClient.Close()
|
|
|
|
redisClient.FlushDB(context.Background())
|
|
|
|
jwtUtil := NewJWTUtil("test-secret", 3600, redisClient)
|
|
ctx := context.Background()
|
|
|
|
userID := int64(456)
|
|
username := "refreshuser"
|
|
|
|
// 生成原始 token
|
|
oldToken, err := jwtUtil.GenerateToken(ctx, userID, username)
|
|
assert.NoError(t, err)
|
|
|
|
// 刷新 token
|
|
newToken, err := jwtUtil.RefreshToken(ctx, oldToken)
|
|
assert.NoError(t, err)
|
|
assert.NotEmpty(t, newToken)
|
|
assert.NotEqual(t, oldToken, newToken)
|
|
|
|
// 验证新 token
|
|
claims, err := jwtUtil.ValidateToken(ctx, newToken)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, userID, claims.UserID)
|
|
assert.Equal(t, username, claims.Username)
|
|
|
|
// 验证旧 token 已失效
|
|
_, err = jwtUtil.ValidateToken(ctx, oldToken)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
// 测试获取和删除用户所有 token
|
|
func TestJWTUtil_UserTokens(t *testing.T) {
|
|
redisClient := setupTestRedis()
|
|
defer redisClient.Close()
|
|
|
|
redisClient.FlushDB(context.Background())
|
|
|
|
jwtUtil := NewJWTUtil("test-secret", 3600, redisClient)
|
|
ctx := context.Background()
|
|
|
|
userID := int64(789)
|
|
username := "multiuser"
|
|
|
|
// 生成多个 token
|
|
var tokens []string
|
|
for i := 0; i < 3; i++ {
|
|
token, err := jwtUtil.GenerateToken(ctx, userID, username)
|
|
assert.NoError(t, err)
|
|
tokens = append(tokens, token)
|
|
}
|
|
|
|
// 获取用户所有 token
|
|
userTokens, err := jwtUtil.GetUserTokens(ctx, userID)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, 3, len(userTokens))
|
|
|
|
// 删除用户所有 token
|
|
err = jwtUtil.DeleteAllUserTokens(ctx, userID)
|
|
assert.NoError(t, err)
|
|
|
|
// 验证所有 token 已被删除
|
|
for _, token := range tokens {
|
|
exists, err := jwtUtil.TokenExists(ctx, token)
|
|
assert.NoError(t, err)
|
|
assert.False(t, exists)
|
|
}
|
|
}
|
|
|
|
// 测试无效 token
|
|
func TestJWTUtil_InvalidToken(t *testing.T) {
|
|
redisClient := setupTestRedis()
|
|
defer redisClient.Close()
|
|
|
|
redisClient.FlushDB(context.Background())
|
|
|
|
jwtUtil := NewJWTUtil("test-secret", 3600, redisClient)
|
|
ctx := context.Background()
|
|
|
|
// 测试无效 token
|
|
_, err := jwtUtil.ValidateToken(ctx, "invalid-token")
|
|
assert.Error(t, err)
|
|
|
|
// 测试空 token
|
|
_, err = jwtUtil.ValidateToken(ctx, "")
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
// 测试过期 token
|
|
func TestJWTUtil_ExpiredToken(t *testing.T) {
|
|
redisClient := setupTestRedis()
|
|
defer redisClient.Close()
|
|
|
|
redisClient.FlushDB(context.Background())
|
|
|
|
// 创建一个很短过期时间的 JWT 工具
|
|
jwtUtil := NewJWTUtil("test-secret", 1, redisClient) // 1秒过期
|
|
ctx := context.Background()
|
|
|
|
userID := int64(999)
|
|
username := "expireduser"
|
|
|
|
// 生成 token
|
|
token, err := jwtUtil.GenerateToken(ctx, userID, username)
|
|
assert.NoError(t, err)
|
|
|
|
// 立即验证应该成功
|
|
claims, err := jwtUtil.ValidateToken(ctx, token)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, userID, claims.UserID)
|
|
|
|
// 等待 token 过期
|
|
time.Sleep(2 * time.Second)
|
|
|
|
// 验证过期 token 应该失败
|
|
_, err = jwtUtil.ValidateToken(ctx, token)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
// 测试不同签名密钥
|
|
func TestJWTUtil_DifferentSecret(t *testing.T) {
|
|
redisClient := setupTestRedis()
|
|
defer redisClient.Close()
|
|
|
|
redisClient.FlushDB(context.Background())
|
|
|
|
jwtUtil1 := NewJWTUtil("secret1", 3600, redisClient)
|
|
jwtUtil2 := NewJWTUtil("secret2", 3600, redisClient)
|
|
ctx := context.Background()
|
|
|
|
userID := int64(111)
|
|
username := "secretuser"
|
|
|
|
// 用第一个工具生成 token
|
|
token, err := jwtUtil1.GenerateToken(ctx, userID, username)
|
|
assert.NoError(t, err)
|
|
|
|
// 用第一个工具验证应该成功
|
|
claims, err := jwtUtil1.ValidateToken(ctx, token)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, userID, claims.UserID)
|
|
|
|
// 用第二个工具验证应该失败
|
|
_, err = jwtUtil2.ValidateToken(ctx, token)
|
|
assert.Error(t, err)
|
|
}
|
|
|
|
// 基准测试：生成 token
|
|
func BenchmarkJWTUtil_GenerateToken(b *testing.B) {
|
|
redisClient := setupTestRedis()
|
|
defer redisClient.Close()
|
|
|
|
redisClient.FlushDB(context.Background())
|
|
|
|
jwtUtil := NewJWTUtil("bench-secret", 3600, redisClient)
|
|
ctx := context.Background()
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
_, err := jwtUtil.GenerateToken(ctx, int64(i), "benchuser")
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// 基准测试：验证 token
|
|
func BenchmarkJWTUtil_ValidateToken(b *testing.B) {
|
|
redisClient := setupTestRedis()
|
|
defer redisClient.Close()
|
|
|
|
redisClient.FlushDB(context.Background())
|
|
|
|
jwtUtil := NewJWTUtil("bench-secret", 3600, redisClient)
|
|
ctx := context.Background()
|
|
|
|
// 预先生成一些 token
|
|
tokens := make([]string, 1000)
|
|
for i := 0; i < 1000; i++ {
|
|
token, err := jwtUtil.GenerateToken(ctx, int64(i), "benchuser")
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
tokens[i] = token
|
|
}
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
tokenIndex := i % len(tokens)
|
|
_, err := jwtUtil.ValidateToken(ctx, tokens[tokenIndex])
|
|
if err != nil {
|
|
b.Fatal(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// 示例：如何在 HTTP 处理器中使用
|
|
func ExampleJWTUtil_HTTPHandler() {
|
|
redisClient := redis.NewClient(&redis.Options{
|
|
Addr: "localhost:6379",
|
|
DB: 0,
|
|
})
|
|
defer redisClient.Close()
|
|
|
|
jwtUtil := NewJWTUtil("your-secret", 3600, redisClient)
|
|
ctx := context.Background()
|
|
|
|
// 模拟登录处理
|
|
handleLogin := func(userID int64, username string) (string, error) {
|
|
token, err := jwtUtil.GenerateToken(ctx, userID, username)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
return token, nil
|
|
}
|
|
|
|
// 模拟需要认证的处理
|
|
handleProtected := func(tokenString string) (*JWTClaims, error) {
|
|
claims, err := jwtUtil.ValidateToken(ctx, tokenString)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return claims, nil
|
|
}
|
|
|
|
// 模拟登出处理
|
|
handleLogout := func(tokenString string) error {
|
|
return jwtUtil.DeleteToken(ctx, tokenString)
|
|
}
|
|
|
|
// 示例使用
|
|
userID := int64(123)
|
|
username := "exampleuser"
|
|
|
|
// 登录
|
|
token, err := handleLogin(userID, username)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
// 访问保护的资源
|
|
claims, err := handleProtected(token)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
// 登出
|
|
err = handleLogout(token)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
// 输出结果
|
|
t := &testing.T{}
|
|
assert.NotEmpty(t, token)
|
|
assert.Equal(t, userID, claims.UserID)
|
|
assert.Equal(t, username, claims.Username)
|
|
}
|
|
|