99 lines
1.8 KiB
Go
99 lines
1.8 KiB
Go
package metrics
|
|
|
|
import (
|
|
"bufio"
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gorilla/mux"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
var dataRow = []byte("I'm the data Row\n")
|
|
|
|
func chunkedHandler(w http.ResponseWriter, r *http.Request) {
|
|
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(dataRow)
|
|
if f, ok := w.(http.Flusher); ok {
|
|
f.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()
|
|
}
|
|
|
|
func TestChunked(t *testing.T) {
|
|
mr := mux.NewRouter()
|
|
mr.Use(HTTPMetricMiddleware)
|
|
mr.Handle("/", http.HandlerFunc(chunkedHandler))
|
|
s := httptest.NewServer(mr)
|
|
defer s.Close()
|
|
|
|
resp, err := http.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, dataRow, append(line, '\n'))
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|