GitOps for k8s
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

websocket_test.go 1.7KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. package websocket
  2. import (
  3. "bytes"
  4. "io"
  5. "net/http"
  6. "net/http/httptest"
  7. "net/url"
  8. "sync"
  9. "testing"
  10. "github.com/fluxcd/flux/http/client"
  11. )
  12. func TestToken(t *testing.T) {
  13. token := "toooookkkkkeeeeennnnnn"
  14. upgrade := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  15. tok := r.Header.Get("Authorization")
  16. if tok != "Scope-Probe token="+token {
  17. t.Fatal("Did not get authorisation header, got: " + tok)
  18. }
  19. _, err := Upgrade(w, r, nil)
  20. if err != nil {
  21. t.Fatal(err)
  22. }
  23. })
  24. srv := httptest.NewServer(upgrade)
  25. defer srv.Close()
  26. url, _ := url.Parse(srv.URL)
  27. url.Scheme = "ws"
  28. ws, err := Dial(http.DefaultClient, "fluxd/test", client.Token(token), url)
  29. if err != nil {
  30. t.Fatal(err)
  31. }
  32. defer ws.Close()
  33. }
  34. func TestByteStream(t *testing.T) {
  35. buf := &bytes.Buffer{}
  36. var wg sync.WaitGroup
  37. wg.Add(1)
  38. upgrade := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  39. ws, err := Upgrade(w, r, nil)
  40. if err != nil {
  41. t.Fatal(err)
  42. }
  43. if _, err := io.Copy(buf, ws); err != nil {
  44. t.Fatal(err)
  45. }
  46. wg.Done()
  47. })
  48. srv := httptest.NewServer(upgrade)
  49. url, _ := url.Parse(srv.URL)
  50. url.Scheme = "ws"
  51. ws, err := Dial(http.DefaultClient, "fluxd/test", client.Token(""), url)
  52. if err != nil {
  53. t.Fatal(err)
  54. }
  55. checkWrite := func(msg string) {
  56. if _, err := ws.Write([]byte(msg)); err != nil {
  57. t.Fatal(err)
  58. }
  59. }
  60. checkWrite("hey")
  61. checkWrite(" there")
  62. checkWrite(" champ")
  63. if err := ws.Close(); err != nil {
  64. t.Fatal(err)
  65. }
  66. // Make sure the server reads everything from the connection
  67. srv.Close()
  68. wg.Wait()
  69. if buf.String() != "hey there champ" {
  70. t.Fatalf("did not collect message as expected, got %s", buf.String())
  71. }
  72. }