blob: b18abc146fe796a9fa957e56a8ed764e8a1b169a [file] [log] [blame]
Tatiana Bradley7c9dd8d2023-08-29 12:16:24 -04001// Copyright 2023 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package proxy
6
7import (
8 "encoding/json"
9 "net/http"
10 "net/http/httptest"
11 "os"
12 "path/filepath"
13 "sync"
14 "testing"
15)
16
17// NewTestClient creates a new client for testing.
18// If update is true, the returned client contacts the real
19// proxy and updates the file "testdata/proxy/<TestName>.json" with
20// the responses it saw.
21// If update is false, the returned client is a fake that
22// reads saved responses from "testdata/proxy/<TestName>.json".
23func NewTestClient(t *testing.T, update bool) (*Client, error) {
24 t.Helper()
25
26 fpath := responsesFile(t)
27 if update {
28 // Set up a real proxy and register a function to write the responses
29 // after the test runs.
30 pc := NewClient(http.DefaultClient, ProxyURL)
31 t.Cleanup(func() {
32 if err := os.MkdirAll(filepath.Dir(fpath), os.ModePerm); err != nil {
33 t.Error(err)
34 return
35 }
36 if err := pc.writeResponses(fpath); err != nil {
37 t.Error(err)
38 }
39 })
40 return pc, nil
41 }
42
43 // Get the fake client from the saved responses.
44 b, err := os.ReadFile(fpath)
45 if err != nil {
46 return nil, err
47 }
48 var responses map[string]*response
49 err = json.Unmarshal(b, &responses)
50 if err != nil {
51 return nil, err
52 }
53 c, cleanup := fakeClient(responses)
54 t.Cleanup(cleanup)
55
56 return c, nil
57}
58
59// response is a representation of an HTTP response used to
60// facilitate testing.
61type response struct {
62 Body string `json:"body,omitempty"`
63 StatusCode int `json:"status_code"`
64}
65
66// fakeClient creates a client that returns hard-coded responses.
67// endpointsToResponses is a map from proxy endpoints
68// (with no server url, and no leading '/'), to their desired responses.
69func fakeClient(endpointsToResponses map[string]*response) (c *Client, cleanup func()) {
70 handler := func(w http.ResponseWriter, r *http.Request) {
71 for endpoint, response := range endpointsToResponses {
72 if r.Method == http.MethodGet &&
73 r.URL.Path == "/"+endpoint {
74 if response.StatusCode == http.StatusOK {
75 _, _ = w.Write([]byte(response.Body))
76 } else {
77 w.WriteHeader(response.StatusCode)
78 }
79 return
80 }
81 }
82 w.WriteHeader(http.StatusBadRequest)
83 }
84 s := httptest.NewServer(http.HandlerFunc(handler))
85 return NewClient(s.Client(), s.URL), func() { s.Close() }
86}
87
88func responsesFile(t *testing.T) string {
89 return filepath.Join("testdata", "proxy", t.Name()+".json")
90}
91
92// responses returns a map from endpoints to the latest response received for each endpoint.
93//
94// Intended for testing: the output can be passed to NewTestClient to create a fake client
95// that returns the same responses.
96func (c *Client) responses() map[string]*response {
97 m := make(map[string]*response)
98 for key, status := range c.errLog.getData() {
99 m[key] = &response{StatusCode: status}
100 }
101 for key, b := range c.cache.getData() {
102 m[key] = &response{Body: string(b), StatusCode: http.StatusOK}
103 }
104 return m
105}
106
107func (pc *Client) writeResponses(filepath string) error {
108 responses, err := json.MarshalIndent(pc.responses(), "", "\t")
109 if err != nil {
110 return err
111 }
112 return os.WriteFile(filepath, responses, 0644)
113}
114
115// An in-memory store of the errors seen so far.
116// Used by the responses() function, for testing.
117type errLog struct {
118 data map[string]int
119 mu sync.Mutex
120}
121
122func newErrLog() *errLog {
123 return &errLog{data: make(map[string]int)}
124}
125
126func (e *errLog) set(key string, status int) {
127 e.mu.Lock()
128 defer e.mu.Unlock()
129
130 e.data[key] = status
131}
132
133func (e *errLog) getData() map[string]int {
134 e.mu.Lock()
135 defer e.mu.Unlock()
136
137 return e.data
138}