| 1885 | } |
| 1886 | |
| 1887 | func TestAutoDrainBodyWithPartialRead(t *testing.T) { |
| 1888 | var actualConnections int32 |
| 1889 | ts := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { |
| 1890 | w.Header().Set("X-Elastic-Product", "Elasticsearch") |
| 1891 | w.WriteHeader(http.StatusOK) |
| 1892 | if _, err := w.Write([]byte(`{"cluster_name":"test","test_data":"` + strings.Repeat("x", 3000) + `"}`)); err != nil { |
| 1893 | t.Errorf("failed to write response: %v", err) |
| 1894 | } |
| 1895 | })) |
| 1896 | ts.Config.ConnState = func(_ net.Conn, state http.ConnState) { |
| 1897 | if state == http.StateNew { |
| 1898 | atomic.AddInt32(&actualConnections, 1) |
| 1899 | } |
| 1900 | } |
| 1901 | ts.Start() |
| 1902 | defer ts.Close() |
| 1903 | es, err := NewClient(Config{ |
| 1904 | Addresses: []string{ts.URL}, |
| 1905 | AutoDrainBody: true, |
| 1906 | }) |
| 1907 | if err != nil { |
| 1908 | t.Fatalf("Failed to create client: %v", err) |
| 1909 | } |
| 1910 | for range 3 { |
| 1911 | res, err := es.Info() |
| 1912 | if err != nil { |
| 1913 | t.Fatalf("Request failed: %v", err) |
| 1914 | } |
| 1915 | buf := make([]byte, 10) |
| 1916 | n, err := res.Body.Read(buf) |
| 1917 | if err != nil && err != io.EOF { |
| 1918 | t.Fatalf("Read error: %v", err) |
| 1919 | } |
| 1920 | if n == 0 { |
| 1921 | t.Error("Expected to read some data") |
| 1922 | } |
| 1923 | if err := res.Body.Close(); err != nil { |
| 1924 | t.Fatalf("Close error: %v", err) |
| 1925 | } |
| 1926 | } |
| 1927 | if atomic.LoadInt32(&actualConnections) > 1 { |
| 1928 | t.Errorf("Expected connection reuse after partial read (1 connection), but saw %d", actualConnections) |
| 1929 | } |
| 1930 | } |