在单元测试中模拟context.Done（）

Question

I have a HTTP handler that sets a context deadline on each request:

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)
        }
    }
}

I am easily able to test the http.StatusNoContent header, but I am unsure about how to test the <-ctx.Done() case in the select statement.

In my test case I have built a mock context.Context and passed it to the req.WithContext() method on my mock http.Request, however, the status code returned is always http.StatusNoContent which leads me to believe the select statement is always falling into the first case in my test.

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)
    }
}

How could I mock the context deadline has exceeded?

Answer 1

So, after much trial and error I figured out what I was doing wrong. Instead of trying to create a mock context.Context, I created a new one with an expired deadline and immediately called the returned cancelFunc. I then passed this to req.WithContext() and now it works like a charm!

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)
    }
}