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

Question

我有一个 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 header ，但我不确定如何测试 <-ctx.Done() select 语句中的大小写。

在我的测试用例中，我构建了一个模拟 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)
    }
}

我如何模拟已超过上下文截止日期？

Answer 1

所以，经过多次尝试和错误，我弄清楚了我做错了什么。我没有尝试创建一个模拟 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)
    }
}