internal/testutil/status_writer.go

   1 package testutil
   2
   3 import (
   4         "fmt"
   5         "io"
   6         "log"
   7         "net"
   8         "net/http"
   9         "testing"
  10
  11         "github.com/anacrolix/missinggo"
  12         "github.com/stretchr/testify/require"
  13 )
  14
  15 type StatusWriter interface {
  16         WriteStatus(io.Writer)
  17 }
  18
  19 // Use StatusServer instead to allow -count > 1 when testing.
  20 func ExportStatusWriter(sw StatusWriter, path string) {
  21         http.HandleFunc(
  22                 fmt.Sprintf("/%s/%s", missinggo.GetTestName(), path),
  23                 func(w http.ResponseWriter, r *http.Request) {
  24                         sw.WriteStatus(w)
  25                 },
  26         )
  27 }
  28
  29 type StatusServer struct {
  30         sm http.ServeMux
  31         l  net.Listener
  32 }
  33
  34 func NewStatusServer(t *testing.T) (ret *StatusServer) {
  35         l, err := net.Listen("tcp", "localhost:0")
  36         require.NoError(t, err)
  37         ret = &StatusServer{
  38                 l: l,
  39         }
  40         log.Printf("serving status at %q", l.Addr())
  41         go http.Serve(l, &ret.sm)
  42         return
  43 }
  44
  45 func (me *StatusServer) HandleStatusWriter(sw StatusWriter, path string) {
  46         me.sm.HandleFunc(path, func(w http.ResponseWriter, r *http.Request) {
  47                 sw.WriteStatus(w)
  48         })
  49 }
  50
  51 func (me StatusServer) Close() {
  52         me.l.Close()
  53 }