|
| 1 | +package gocksnap |
| 2 | + |
| 3 | +import ( |
| 4 | + _ "embed" |
| 5 | + "encoding/json" |
| 6 | + "fmt" |
| 7 | + "io" |
| 8 | + "net/http" |
| 9 | + "os" |
| 10 | + "path/filepath" |
| 11 | + "strings" |
| 12 | + "sync" |
| 13 | + "testing" |
| 14 | + |
| 15 | + "github.com/h2non/gock" |
| 16 | +) |
| 17 | + |
| 18 | +const defaultSnapshotDirectory = "__snapshots__" |
| 19 | + |
| 20 | +//go:embed index.html |
| 21 | +var indexHTML string |
| 22 | + |
| 23 | +// Call represents a single HTTP call in the snapshot. |
| 24 | +type Call struct { |
| 25 | + // Method is the HTTP method of the call (GET, POST, etc.) |
| 26 | + Method string `json:"method"` |
| 27 | + |
| 28 | + // URL is the full URL of the call. |
| 29 | + URL string `json:"url"` |
| 30 | + |
| 31 | + // ReqBody is the request body of the call, if any. |
| 32 | + ReqBody json.RawMessage `json:"reqBody"` |
| 33 | + |
| 34 | + // ResBody is the response body of the call, if any. |
| 35 | + ResBody json.RawMessage `json:"resBody"` |
| 36 | + |
| 37 | + // Status is the HTTP status code of the response. |
| 38 | + Status int `json:"status"` |
| 39 | +} |
| 40 | + |
| 41 | +// Snapshot holds the state of the snapshot being recorded, which can include multiple HTTP calls. |
| 42 | +type Snapshot struct { |
| 43 | + Calls []Call `json:"calls"` |
| 44 | + |
| 45 | + // testName used to identify the snapshot file. |
| 46 | + testName string |
| 47 | + |
| 48 | + // name is the name of the snapshot, used for identification in the test and in the file. |
| 49 | + name string |
| 50 | + |
| 51 | + // updateMode indicates if the snapshot is in update mode or in test mode. |
| 52 | + updateMode bool |
| 53 | + |
| 54 | + // mu is a mutex to protect access to the pending call and SSE connections. |
| 55 | + mu sync.Mutex |
| 56 | + |
| 57 | + // pending is the current call that is being recorded. |
| 58 | + pending *CallPrompt |
| 59 | + |
| 60 | + // sseConns is a map of client connections that are waiting for updates. |
| 61 | + sseConns map[chan string]struct{} |
| 62 | +} |
| 63 | + |
| 64 | +// Finish |
| 65 | +func (g *Snapshot) Finish(t *testing.T) { |
| 66 | + t.Helper() |
| 67 | + |
| 68 | + if !g.updateMode { |
| 69 | + if !gock.IsDone() { |
| 70 | + t.Fatalf("Snapshot '%s' is not complete. Some requests were not mocked.", g.name) |
| 71 | + } |
| 72 | + |
| 73 | + return |
| 74 | + } |
| 75 | + |
| 76 | + // Update snapshot |
| 77 | + data, err := json.MarshalIndent(g, "", " ") |
| 78 | + if err != nil { |
| 79 | + t.Fatalf("Failed to marshal snapshot '%s': %v", g.name, err) |
| 80 | + } |
| 81 | + |
| 82 | + _ = os.MkdirAll(defaultSnapshotDirectory, 0o750) |
| 83 | + |
| 84 | + err = os.WriteFile(g.file(), data, 0o600) |
| 85 | + if err != nil { |
| 86 | + t.Fatalf("Failed to save snapshot '%s': %v", g.name, err) |
| 87 | + } |
| 88 | +} |
| 89 | + |
| 90 | +// file returns the path to the snapshot file. |
| 91 | +func (g *Snapshot) file() string { |
| 92 | + return filepath.Join(defaultSnapshotDirectory, strings.ReplaceAll(strings.ReplaceAll(g.testName+"-"+g.name, " ", "_"), "/", "_")+".json") |
| 93 | +} |
| 94 | + |
| 95 | +// promptCall sends the current request to the UI for user interaction. |
| 96 | +func (g *Snapshot) promptCall(req *http.Request, existingCall *Call) *Call { |
| 97 | + var bodyRaw []byte |
| 98 | + if req.Body != nil { |
| 99 | + bodyRaw, _ = io.ReadAll(req.Body) |
| 100 | + } |
| 101 | + |
| 102 | + g.mu.Lock() |
| 103 | + |
| 104 | + fmt.Printf("Request: %s %s\n", req.Method, req.URL.String()) |
| 105 | + |
| 106 | + g.pending = &CallPrompt{ |
| 107 | + Name: g.name, |
| 108 | + Call: Call{ |
| 109 | + Method: req.Method, |
| 110 | + URL: req.URL.String(), |
| 111 | + ReqBody: bodyRaw, |
| 112 | + }, |
| 113 | + ExistingCall: existingCall, |
| 114 | + finalCall: make(chan *Call, 1), |
| 115 | + } |
| 116 | + |
| 117 | + // notify SSE clients |
| 118 | + for ch := range g.sseConns { |
| 119 | + select { |
| 120 | + case ch <- "pending": |
| 121 | + default: |
| 122 | + } |
| 123 | + } |
| 124 | + |
| 125 | + g.mu.Unlock() |
| 126 | + |
| 127 | + return <-g.pending.finalCall |
| 128 | +} |
| 129 | + |
| 130 | +// MatchSnapshot creates a new snapshot for the current test. |
| 131 | +// If the snapshot file is not found, or if the environment variable UPDATE_TESTS is set to "true", it will spawn a web server to allow the user to interactively select responses for the recorded requests. |
| 132 | +// If the snapshot file is found, it will load the existing calls and register them with gock. |
| 133 | +// After all the calls are finished, the user should call the Finish method to save the snapshot / assert that all calls were mocked correctly. |
| 134 | +func MatchSnapshot(t *testing.T, snapshotName string) *Snapshot { |
| 135 | + t.Helper() |
| 136 | + |
| 137 | + snapshot := &Snapshot{ |
| 138 | + Calls: []Call{}, |
| 139 | + testName: t.Name(), |
| 140 | + name: snapshotName, |
| 141 | + updateMode: os.Getenv("UPDATE_TESTS") == "true", |
| 142 | + sseConns: make(map[chan string]struct{}), |
| 143 | + } |
| 144 | + |
| 145 | + var existingCalls []Call |
| 146 | + |
| 147 | + _, err := os.Stat(snapshot.file()) |
| 148 | + if os.IsNotExist(err) { |
| 149 | + // can't find snapshot file, so we are in update mode |
| 150 | + t.Logf("Snapshot '%s' not found, running in update mode\n", snapshot.file()) |
| 151 | + snapshot.updateMode = true |
| 152 | + } else { |
| 153 | + // Load existing snapshot |
| 154 | + data, err := os.ReadFile(snapshot.file()) |
| 155 | + if err != nil { |
| 156 | + t.Fatalf("Failed to open snapshot '%s': %v", snapshot.file(), err) |
| 157 | + } |
| 158 | + |
| 159 | + err = json.Unmarshal(data, snapshot) |
| 160 | + if err != nil { |
| 161 | + t.Fatalf("Failed to unmarshal snapshot '%s': %v", snapshot.file(), err) |
| 162 | + } |
| 163 | + |
| 164 | + if snapshot.updateMode { |
| 165 | + existingCalls = snapshot.Calls |
| 166 | + snapshot.Calls = make([]Call, 0) |
| 167 | + t.Logf("Updating existing snapshot '%s'\n", snapshot.file()) |
| 168 | + } |
| 169 | + } |
| 170 | + |
| 171 | + if snapshot.updateMode { |
| 172 | + addr, err := snapshot.startPromptServer() |
| 173 | + if err != nil { |
| 174 | + t.Fatalf("Failed to start prompt server for snapshot '%s': %v", snapshot.file(), err) |
| 175 | + } |
| 176 | + |
| 177 | + openBrowser(addr) |
| 178 | + } |
| 179 | + |
| 180 | + gock.Intercept() |
| 181 | + |
| 182 | + if snapshot.updateMode { |
| 183 | + var existingCall *Call |
| 184 | + if len(existingCalls) > 0 { |
| 185 | + existingCall = &existingCalls[0] |
| 186 | + } |
| 187 | + |
| 188 | + // clean up any existing mocks |
| 189 | + for _, mock := range gock.Pending() { |
| 190 | + mock.Disable() |
| 191 | + } |
| 192 | + |
| 193 | + gock.Register(snapshot.newRecordMock(existingCall)) |
| 194 | + gock.Observe(func(_ *http.Request, mock gock.Mock) { |
| 195 | + snapshot.Calls = append(snapshot.Calls, *mock.(*recordMocker).call) |
| 196 | + // load the next one |
| 197 | + existingCall = nil |
| 198 | + if len(snapshot.Calls) < len(existingCalls) { |
| 199 | + existingCall = &existingCalls[len(snapshot.Calls)] |
| 200 | + } |
| 201 | + |
| 202 | + gock.Register(snapshot.newRecordMock(existingCall)) |
| 203 | + }) |
| 204 | + |
| 205 | + return snapshot |
| 206 | + } |
| 207 | + |
| 208 | + // Register existing calls into gock. |
| 209 | + for _, call := range snapshot.Calls { |
| 210 | + req := gock.NewRequest().URL(call.URL).JSON(call.ReqBody) |
| 211 | + req.Method = strings.ToUpper(call.Method) |
| 212 | + gock.Register(gock.NewMock(req, gock.NewResponse().Status(call.Status).JSON(call.ResBody))) |
| 213 | + } |
| 214 | + |
| 215 | + t.Logf("Loaded snapshot '%s' with %d calls", snapshot.file(), len(snapshot.Calls)) |
| 216 | + |
| 217 | + return snapshot |
| 218 | +} |
0 commit comments