internal/testutil/status_writer.go

   1 package testutil
   2
   3 import (
   4         "fmt"
   5         "io"
   6         "net/http"
   7         "sync"
   8         "testing"
   9
  10         _ "github.com/anacrolix/envpprof"
  11 )
  12
  13 type StatusWriter interface {
  14         WriteStatus(io.Writer)
  15 }
  16
  17 // The key is the route pattern. The value is nil when the resource is released.
  18 var (
  19         mu  sync.Mutex
  20         sws = map[string]StatusWriter{}
  21 )
  22
  23 func ExportStatusWriter(sw StatusWriter, path string, t testing.TB) (release func()) {
  24         pattern := fmt.Sprintf("/%s/%s", t.Name(), path)
  25         t.Logf("exporting status path %q", pattern)
  26         release = func() {
  27                 mu.Lock()
  28                 defer mu.Unlock()
  29                 sws[pattern] = nil
  30         }
  31         mu.Lock()
  32         defer mu.Unlock()
  33         if curSw, ok := sws[pattern]; ok {
  34                 if curSw != nil {
  35                         panic(fmt.Sprintf("%q still in use", pattern))
  36                 }
  37                 sws[pattern] = sw
  38                 return
  39         }
  40         http.HandleFunc(
  41                 pattern,
  42                 func(w http.ResponseWriter, r *http.Request) {
  43                         sw := sws[pattern]
  44                         if sw == nil {
  45                                 http.NotFound(w, r)
  46                                 return
  47                         }
  48                         sw.WriteStatus(w)
  49                 },
  50         )
  51         sws[pattern] = sw
  52         return
  53 }