[Go] Stream HTTP Response with chunked transfer encoding.

chunked_test.go

package main

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/http/httptest"
	"sync"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"

	"github.com/stretchr/testify/require"
)

var lineDelimiter = []byte("\n") // line delimiter

var payload = []byte(`I'm the final response!`)

func TestStreamResponse(t *testing.T) {
	s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Content-Type-Options", "nosniff")

		ctx, cancel := context.WithCancel(r.Context())
		ticker := time.NewTicker(time.Second) // We may set it to 10 secs
		var wg sync.WaitGroup
		wg.Add(1)
		go func() {
			defer wg.Done()
			for {
				select {
				case <-ticker.C:
					_, _ = w.Write(lineDelimiter)
					w.(http.Flusher).Flush()
				case <-ctx.Done():
					return
				}
			}
		}()

		// Emulate some work
		time.Sleep(5 * time.Second)

		// Telling the loop that keeps the connection alive to end
		cancel()

		// Waiting until the loop ends
		wg.Wait()

		// Sending back the final response
		_, _ = w.Write(payload)

	}))

	defer s.Close()

	c := http.Client{
		Timeout: time.Second * 3, // Remove if want the test to SUCCESS
	}

	resp, err := c.Get(s.URL)
	require.NoError(t, err)
	assert.Contains(t, resp.TransferEncoding, "chunked")

	defer resp.Body.Close()

	r := bufio.NewReader(resp.Body)
	for {
		line, err := readChunkedResponseLine(r)
		if err != nil {
			if err == io.EOF {
				return
			}
			log.Fatal(err.Error())
		}
		if len(line) == 0 {
			log.Println("Alive!")
			continue
		}

		fmt.Println(string(line)) // we got the final response
		assert.Equal(t, payload, line)
	}
}

func readChunkedResponseLine(r *bufio.Reader) ([]byte, error) {
	line, isPrefix, err := r.ReadLine()
	if err != nil {
		return nil, err
	}

	if isPrefix {
		rest, err := readChunkedResponseLine(r)
		if err != nil {
			return nil, err
		}
		line = append(line, rest...)
	}

	return line, nil
}