| 20 | ) |
| 21 | |
| 22 | func Test_ProxyServer_Headers(t *testing.T) { |
| 23 | t.Parallel() |
| 24 | |
| 25 | const ( |
| 26 | headerName1 = "X-Test-Header-1" |
| 27 | headerVal1 = "test-value-1" |
| 28 | headerName2 = "X-Test-Header-2" |
| 29 | headerVal2 = "test-value-2" |
| 30 | ) |
| 31 | |
| 32 | // We're not going to actually start a proxy, we're going to point it |
| 33 | // towards a fake server that returns an unexpected status code. This'll |
| 34 | // cause the proxy to exit with an error that we can check for. |
| 35 | var called atomic.Int64 |
| 36 | srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 37 | called.Add(1) |
| 38 | assert.Equal(t, headerVal1, r.Header.Get(headerName1)) |
| 39 | assert.Equal(t, headerVal2, r.Header.Get(headerName2)) |
| 40 | |
| 41 | w.WriteHeader(http.StatusTeapot) // lol |
| 42 | })) |
| 43 | defer srv.Close() |
| 44 | |
| 45 | inv, _ := newCLI(t, "wsproxy", "server", |
| 46 | "--primary-access-url", srv.URL, |
| 47 | "--proxy-session-token", "test-token", |
| 48 | "--access-url", "http://localhost:8080", |
| 49 | "--http-address", ":0", |
| 50 | "--header", fmt.Sprintf("%s=%s", headerName1, headerVal1), |
| 51 | "--header-command", fmt.Sprintf("printf %s=%s", headerName2, headerVal2), |
| 52 | ) |
| 53 | pty := ptytest.New(t) |
| 54 | inv.Stdout = pty.Output() |
| 55 | err := inv.Run() |
| 56 | require.Error(t, err) |
| 57 | require.ErrorContains(t, err, "unexpected status code 418") |
| 58 | require.NoError(t, pty.Close()) |
| 59 | |
| 60 | assert.EqualValues(t, 1, called.Load()) |
| 61 | } |
| 62 | |
| 63 | //nolint:paralleltest,tparallel // Test uses a static port. |
| 64 | func TestWorkspaceProxy_Server_PrometheusEnabled(t *testing.T) { |