在单元测试中模拟context.Done()

时间:2017-05-12 21:55:24

标签: unit-testing go

我有一个HTTP处理程序,可以为每个请求设置上下文截止日期:

func submitHandler(stream chan data) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
        defer cancel()

        // read request body, etc.

        select {
        case stream <- req:
            w.WriteHeader(http.StatusNoContent)
        case <-ctx.Done():
            err := ctx.Err()
            if err == context.DeadlineExceeded {
                w.WriteHeader(http.StatusRequestTimeout)
            }
            log.Printf("context done: %v", err)
        }
    }
}

我可以轻松测试http.StatusNoContent标题,但我不确定如何在select语句中测试<-ctx.Done()个案。

在我的测试用例中,我构建了一个模拟context.Context并将其传递给模拟req.WithContext()上的http.Request方法,但是,返回的状态代码始终为http.StatusNoContent这让我相信select语句总是落入我测试的第一个案例中。

type mockContext struct{}

func (ctx mockContext) Deadline() (deadline time.Time, ok bool) {
    return deadline, ok
}

func (ctx mockContext) Done() <-chan struct{} {
    ch := make(chan struct{})
    close(ch)
    return ch
}

func (ctx mockContext) Err() error {
    return context.DeadlineExceeded
}

func (ctx mockContext) Value(key interface{}) interface{} {
    return nil
}

func TestHandler(t *testing.T) {
    stream := make(chan data, 1)
    defer close(stream)

    handler := submitHandler(stream)
    req, err := http.NewRequest(http.MethodPost, "/submit", nil)
    if err != nil {
        t.Fatal(err)
    }
    req = req.WithContext(mockContext{})

    rec := httptest.NewRecorder()
    handler.ServeHTTP(rec, req)

    if rec.Code != http.StatusRequestTimeout {
        t.Errorf("expected status code: %d, got: %d", http.StatusRequestTimeout, rec.Code)
    }
}

我怎么能模仿上下文截止日期已超过?

2 个答案:

答案 0 :(得分:3)

所以,经过多次试验和错误,我弄明白我做错了什么。我没有尝试创建模拟context.Context,而是创建了一个具有过期期限的新模式,并立即调用返回的cancelFunc。然后我将其传递给req.WithContext(),现在它就像一个魅力!

func TestHandler(t *testing.T) {
    stream := make(chan data, 1)
    defer close(stream)

    handler := submitHandler(stream)
    req, err := http.NewRequest(http.MethodPost, "/submit", nil)
    if err != nil {
        t.Fatal(err)
    }

    stream <- data{}
    ctx, cancel := context.WithDeadline(req.Context(), time.Now().Add(-7*time.Hour))
    cancel()
    req = req.WithContext(ctx)

    rec := httptest.NewRecorder()
    handler.ServeHTTP(rec, req)

    if rec.Code != http.StatusRequestTimeout {
        t.Errorf("expected status code: %d, got: %d", http.StatusRequestTimeout, rec.Code)
    }
}

答案 1 :(得分:1)

您的mockContext类型的Done方法永远不会返回完成,因为没有任何内容写入通道,因此您的goroutine在关闭通道之前会永远坐在那里,从而触发完成状态。如果您希望立即报告完成,请尝试以下操作:

func (ctx mockContext) Done() <-chan struct{} {
    ch := make(chan struct{})
    close(ch)
    return ch
}