Go语言中如何测试AWS Lambda Handler的中间件封装实现

前置说明 #

原有middleware.go代码缺少context包导入，需要先补全导入语句才能正常编译：

import (
    "context" // 补全该导入
    "net/http"

    "github.com/aws/aws-lambda-go/events"
)

测试方案一：依赖注入改造（无第三方依赖、推荐） #

无需修改原有业务调用逻辑，仅新增测试用的兼容方法，即可实现调用次数统计。

第一步：兼容式改造中间件代码 #

修改middleware.go，新增支持注入自定义认证中间件的方法：

// 原有方法保留，完全兼容现有业务调用
func NeedsAdditionalAuth(next middleware) middleware {
    // 内部调用可注入版本，默认传入原有AdditionalAuth实现
    return NeedsAdditionalAuthWithCustomAuth(next, AdditionalAuth)
}

// 新增可注入版本，仅用于单元测试
func NeedsAdditionalAuthWithCustomAuth(next middleware, authMiddleware func(middleware) middleware) middleware {
    return func(ctx context.Context, proxyRequest events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
        if proxyRequest.QueryStringParameters["status"] == "sensitive" {
            return authMiddleware(next)(ctx, proxyRequest)
        }
        return next(ctx, proxyRequest)
    }
}

第二步：编写测试用例 #

新建middleware_test.go，实现调用次数校验逻辑：

package middleware

import (
    "context"
    "net/http"
    "testing"

    "github.com/aws/aws-lambda-go/events"
    "github.com/stretchr/testify/assert"
)

// 测试敏感请求下AdditionalAuth被调用1次
func TestNeedsAdditionalAuth_CallAuthWhenSensitive(t *testing.T) {
    // 初始化调用计数器
    authCallCount := 0
    // 构造Mock认证中间件，统计调用次数
    mockAuth := func(next middleware) middleware {
        return func(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
            authCallCount++
            return next(ctx, req)
        }
    }

    // 构造Mock业务Handler
    handlerCallCount := 0
    mockHandler := func(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
        handlerCallCount++
        return events.APIGatewayProxyResponse{StatusCode: http.StatusOK}, nil
    }

    // 构造敏感请求
    req := events.APIGatewayProxyRequest{
        QueryStringParameters: map[string]string{
            "status": "sensitive",
        },
    }

    // 执行测试
    testChain := NeedsAdditionalAuthWithCustomAuth(mockHandler, mockAuth)
    _, err := testChain(context.Background(), req)

    // 结果断言
    assert.NoError(t, err)
    assert.Equal(t, 1, authCallCount, "敏感请求下AdditionalAuth应当被调用1次")
    assert.Equal(t, 1, handlerCallCount, "业务Handler应当被调用1次")
}

// 测试普通请求下AdditionalAuth不会被调用
func TestNeedsAdditionalAuth_SkipAuthWhenNormal(t *testing.T) {
    authCallCount := 0
    mockAuth := func(next middleware) middleware {
        return func(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
            authCallCount++
            return next(ctx, req)
        }
    }

    handlerCallCount := 0
    mockHandler := func(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
        handlerCallCount++
        return events.APIGatewayProxyResponse{StatusCode: http.StatusOK}, nil
    }

    // 构造普通请求
    req := events.APIGatewayProxyRequest{
        QueryStringParameters: map[string]string{
            "status": "normal",
        },
    }

    testChain := NeedsAdditionalAuthWithCustomAuth(mockHandler, mockAuth)
    _, err := testChain(context.Background(), req)

    assert.NoError(t, err)
    assert.Equal(t, 0, authCallCount, "普通请求下AdditionalAuth不应当被调用")
    assert.Equal(t, 1, handlerCallCount, "业务Handler应当被调用1次")
}

测试方案二：猴子补丁（无需修改原有业务代码） #

如果不想改动业务代码，可以通过猴子补丁替换原有函数实现调用统计，需要引入bouk/monkey库：

func TestNeedsAdditionalAuth_WithMonkeyPatch(t *testing.T) {
    authCallCount := 0
    // 替换原有AdditionalAuth函数
    patch := monkey.Patch(AdditionalAuth, func(next middleware) middleware {
        authCallCount++
        return func(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
            return next(ctx, req)
        }
    })
    defer patch.Unpatch() // 测试结束恢复原函数

    mockHandler := func(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
        return events.APIGatewayProxyResponse{StatusCode: http.StatusOK}, nil
    }

    req := events.APIGatewayProxyRequest{
        QueryStringParameters: map[string]string{"status": "sensitive"},
    }

    testChain := NeedsAdditionalAuth(mockHandler)
    _, err := testChain(context.Background(), req)

    assert.NoError(t, err)
    assert.Equal(t, 1, authCallCount)
}

额外：AdditionalAuth独立逻辑测试 #

如果需要测试认证中间件本身的逻辑，可以用同样的方式MockdoAdditionalAuth的返回值，验证认证通过/失败的分支逻辑：

func TestAdditionalAuth_ForbiddenWhenUnauthenticated(t *testing.T) {
    // Mock认证失败
    patch := monkey.Patch(doAdditionalAuth, func() bool {
        return false
    })
    defer patch.Unpatch()

    handlerCallCount := 0
    mockHandler := func(ctx context.Context, req events.APIGatewayProxyRequest) (events.APIGatewayProxyResponse, error) {
        handlerCallCount++
        return events.APIGatewayProxyResponse{StatusCode: http.StatusOK}, nil
    }

    req := events.APIGatewayProxyRequest{}
    resp, err := AdditionalAuth(mockHandler)(context.Background(), req)

    assert.NoError(t, err)
    assert.Equal(t, http.StatusForbidden, resp.StatusCode)
    assert.Equal(t, 0, handlerCallCount, "认证失败时不应当调用业务Handler")
}

内容的提问来源于stack exchange，提问作者Vi5tar