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