http_service.go 4.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187
  1. // Copyright 2020-2021 InfluxData, Inc. All rights reserved.
  2. // Use of this source code is governed by MIT
  3. // license that can be found in the LICENSE file.
  4. // Package test provides shared test utils
  5. package test
  6. import (
  7. "compress/gzip"
  8. "context"
  9. "fmt"
  10. "io"
  11. "io/ioutil"
  12. "net/http"
  13. "strings"
  14. "sync"
  15. "testing"
  16. http2 "github.com/influxdata/influxdb-client-go/v2/api/http"
  17. "github.com/stretchr/testify/assert"
  18. )
  19. // HTTPService is http.Service implementation for tests
  20. type HTTPService struct {
  21. serverURL string
  22. authorization string
  23. lines []string
  24. t *testing.T
  25. wasGzip bool
  26. requestHandler func(url string, body io.Reader) error
  27. replyError *http2.Error
  28. lock sync.Mutex
  29. requests int
  30. }
  31. // WasGzip returns true of request was in GZip format
  32. func (t *HTTPService) WasGzip() bool {
  33. return t.wasGzip
  34. }
  35. // SetWasGzip sets wasGzip flag
  36. func (t *HTTPService) SetWasGzip(wasGzip bool) {
  37. t.wasGzip = wasGzip
  38. }
  39. // SetRequestHandler sets custom handler for requests
  40. func (t *HTTPService) SetRequestHandler(fn func(url string, body io.Reader) error) {
  41. t.requestHandler = fn
  42. }
  43. // ServerURL returns testing URL
  44. func (t *HTTPService) ServerURL() string {
  45. return t.serverURL
  46. }
  47. // ServerAPIURL returns testing URL
  48. func (t *HTTPService) ServerAPIURL() string {
  49. return t.serverURL
  50. }
  51. // Authorization returns current authorization header value
  52. func (t *HTTPService) Authorization() string {
  53. return t.authorization
  54. }
  55. // HTTPClient returns nil for this service
  56. func (t *HTTPService) HTTPClient() *http.Client {
  57. return nil
  58. }
  59. // Requests returns number of requests
  60. func (t *HTTPService) Requests() int {
  61. return t.requests
  62. }
  63. // Close clears instance
  64. func (t *HTTPService) Close() {
  65. t.lock.Lock()
  66. if len(t.lines) > 0 {
  67. t.lines = t.lines[:0]
  68. }
  69. t.wasGzip = false
  70. t.replyError = nil
  71. t.requestHandler = nil
  72. t.requests = 0
  73. t.lock.Unlock()
  74. }
  75. // SetReplyError sets Error that will be returned as a response
  76. func (t *HTTPService) SetReplyError(replyError *http2.Error) {
  77. t.lock.Lock()
  78. defer t.lock.Unlock()
  79. t.replyError = replyError
  80. }
  81. // ReplyError returns current reply error
  82. func (t *HTTPService) ReplyError() *http2.Error {
  83. t.lock.Lock()
  84. defer t.lock.Unlock()
  85. return t.replyError
  86. }
  87. // SetAuthorization sets authorization string
  88. func (t *HTTPService) SetAuthorization(_ string) {
  89. }
  90. // GetRequest does nothing for this service
  91. func (t *HTTPService) GetRequest(_ context.Context, _ string, _ http2.RequestCallback, _ http2.ResponseCallback) *http2.Error {
  92. return nil
  93. }
  94. // DoHTTPRequest does nothing for this service
  95. func (t *HTTPService) DoHTTPRequest(_ *http.Request, _ http2.RequestCallback, _ http2.ResponseCallback) *http2.Error {
  96. return nil
  97. }
  98. // DoHTTPRequestWithResponse does nothing for this service
  99. func (t *HTTPService) DoHTTPRequestWithResponse(_ *http.Request, _ http2.RequestCallback) (*http.Response, error) {
  100. return nil, nil
  101. }
  102. // DoPostRequest reads http request, validates URL and stores data in the request
  103. func (t *HTTPService) DoPostRequest(_ context.Context, url string, body io.Reader, requestCallback http2.RequestCallback, _ http2.ResponseCallback) *http2.Error {
  104. req, err := http.NewRequest("POST", url, nil)
  105. t.lock.Lock()
  106. t.requests++
  107. t.lock.Unlock()
  108. if err != nil {
  109. return http2.NewError(err)
  110. }
  111. if requestCallback != nil {
  112. requestCallback(req)
  113. }
  114. if req.Header.Get("Content-Encoding") == "gzip" {
  115. body, _ = gzip.NewReader(body)
  116. t.wasGzip = true
  117. }
  118. if t.t != nil {
  119. assert.Equal(t.t, fmt.Sprintf("%swrite?bucket=my-bucket&org=my-org&precision=ns", t.serverURL), url)
  120. }
  121. if t.ReplyError() != nil {
  122. return t.ReplyError()
  123. }
  124. if t.requestHandler != nil {
  125. err = t.requestHandler(url, body)
  126. } else {
  127. err = t.DecodeLines(body)
  128. }
  129. if err != nil {
  130. return http2.NewError(err)
  131. }
  132. return nil
  133. }
  134. // DecodeLines parses request body for lines
  135. func (t *HTTPService) DecodeLines(body io.Reader) error {
  136. bytes, err := ioutil.ReadAll(body)
  137. if err != nil {
  138. return err
  139. }
  140. lines := strings.Split(string(bytes), "\n")
  141. if lines[len(lines)-1] == "" {
  142. lines = lines[:len(lines)-1]
  143. }
  144. t.lock.Lock()
  145. t.lines = append(t.lines, lines...)
  146. t.lock.Unlock()
  147. return nil
  148. }
  149. // Lines returns decoded lines from request
  150. func (t *HTTPService) Lines() []string {
  151. t.lock.Lock()
  152. defer t.lock.Unlock()
  153. return t.lines
  154. }
  155. // NewTestService creates new test HTTP service
  156. func NewTestService(t *testing.T, serverURL string) *HTTPService {
  157. return &HTTPService{
  158. t: t,
  159. serverURL: serverURL + "/api/v2/",
  160. }
  161. }