ooni-probe-cli/internal/bytecounter/http_test.go

package bytecounter

import (
	"context"
	"errors"
	"io"
	"net/http"
	"strings"
	"testing"

	"github.com/ooni/probe-cli/v3/internal/model/mocks"
	"github.com/ooni/probe-cli/v3/internal/netxlite"
)

func TestMaybeWrapHTTPTransport(t *testing.T) {
	t.Run("when counter is not nil", func(t *testing.T) {
		underlying := &mocks.HTTPTransport{}
		counter := &Counter{}
		txp := counter.MaybeWrapHTTPTransport(underlying)
		realTxp := txp.(*httpTransport)
		if realTxp.HTTPTransport != underlying {
			t.Fatal("did not wrap correctly")
		}
	})

	t.Run("when counter is nil", func(t *testing.T) {
		underlying := &mocks.HTTPTransport{}
		var counter *Counter
		txp := counter.MaybeWrapHTTPTransport(underlying)
		if txp != underlying {
			t.Fatal("unexpected result")
		}
	})
}

func TestHTTPTransport(t *testing.T) {
	t.Run("RoundTrip", func(t *testing.T) {
		t.Run("failure", func(t *testing.T) {
			counter := New()
			txp := &httpTransport{
				Counter: counter,
				HTTPTransport: &mocks.HTTPTransport{
					MockRoundTrip: func(req *http.Request) (*http.Response, error) {
						return nil, io.EOF
					},
				},
			}
			req, err := http.NewRequest(
				"POST", "https://www.google.com", strings.NewReader("AAAAAA"))
			if err != nil {
				t.Fatal(err)
			}
			req.Header.Set("User-Agent", "antani-browser/1.0.0")
			resp, err := txp.RoundTrip(req)
			if !errors.Is(err, io.EOF) {
				t.Fatal("not the error we expected")
			}
			if resp != nil {
				t.Fatal("expected nil response here")
			}
			if counter.Sent.Load() != 62 {
				t.Fatal("expected 62 bytes sent", counter.Sent.Load())
			}
			if counter.Received.Load() != 0 {
				t.Fatal("expected zero bytes received", counter.Received.Load())
			}
		})

		t.Run("success", func(t *testing.T) {
			counter := New()
			txp := &httpTransport{
				Counter: counter,
				HTTPTransport: &mocks.HTTPTransport{
					MockRoundTrip: func(req *http.Request) (*http.Response, error) {
						resp := &http.Response{
							Body: io.NopCloser(strings.NewReader("1234567")),
							Header: http.Header{
								"Server": []string{"antani/0.1.0"},
							},
							Status:     "200 OK",
							StatusCode: http.StatusOK,
						}
						return resp, nil
					},
				},
			}
			req, err := http.NewRequest(
				"POST", "https://www.google.com", strings.NewReader("AAAAAA"))
			if err != nil {
				t.Fatal(err)
			}
			req.Header.Set("User-Agent", "antani-browser/1.0.0")
			resp, err := txp.RoundTrip(req)
			if err != nil {
				t.Fatal(err)
			}
			data, err := netxlite.ReadAllContext(context.Background(), resp.Body)
			if err != nil {
				t.Fatal(err)
			}
			resp.Body.Close()
			if string(data) != "1234567" {
				t.Fatal("expected a different body here")
			}
			if counter.Sent.Load() != 62 {
				t.Fatal("expected 62 bytes sent", counter.Sent.Load())
			}
			if counter.Received.Load() != 37 {
				t.Fatal("expected 37 bytes received", counter.Received.Load())
			}
		})

		t.Run("success with EOF", func(t *testing.T) {
			counter := New()
			txp := &httpTransport{
				Counter: counter,
				HTTPTransport: &mocks.HTTPTransport{
					MockRoundTrip: func(req *http.Request) (*http.Response, error) {
						resp := &http.Response{
							Body: io.NopCloser(&mocks.Reader{
								MockRead: func(b []byte) (int, error) {
									if len(b) < 1 {
										panic("should not happen")
									}
									b[0] = 'A'
									return 1, io.EOF // we want code to be robust to this
								},
							}),
							Header: http.Header{
								"Server": []string{"antani/0.1.0"},
							},
							Status:     "200 OK",
							StatusCode: http.StatusOK,
						}
						return resp, nil
					},
				},
			}
			client := &http.Client{Transport: txp}
			resp, err := client.Get("https://www.google.com")
			if err != nil {
				t.Fatal(err)
			}
			data, err := netxlite.ReadAllContext(context.Background(), resp.Body)
			if err != nil {
				t.Fatal(err)
			}
			resp.Body.Close()
			if string(data) != "A" {
				t.Fatal("expected a different body here")
			}
		})
	})

	t.Run("CloseIdleConnections", func(t *testing.T) {
		var called bool
		child := &mocks.HTTPTransport{
			MockCloseIdleConnections: func() {
				called = true
			},
		}
		counter := New()
		txp := WrapHTTPTransport(child, counter)
		txp.CloseIdleConnections()
		if !called {
			t.Fatal("not called")
		}
	})

	t.Run("Network", func(t *testing.T) {
		expected := "antani"
		child := &mocks.HTTPTransport{
			MockNetwork: func() string {
				return expected
			},
		}
		counter := New()
		txp := WrapHTTPTransport(child, counter)
		if network := txp.Network(); network != expected {
			t.Fatal("unexpected network", network)
		}
	})
}