kopia/internal/throttle/round_tripper_test.go

package throttle

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"testing"
)

type baseRoundTripper struct {
	responses map[*http.Request]*http.Response
}

func (rt *baseRoundTripper) add(req *http.Request, resp *http.Response) (*http.Request, *http.Response) {
	rt.responses[req] = resp
	return req, resp
}

func (rt *baseRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	resp := rt.responses[req]
	if resp != nil {
		return resp, nil
	}

	return nil, fmt.Errorf("error occurred")
}

type fakePool struct {
	readers []io.ReadCloser
}

func (fp *fakePool) reset() {
	fp.readers = nil
}

func (fp *fakePool) AddReader(r io.ReadCloser) (io.ReadCloser, error) {
	fp.readers = append(fp.readers, r)
	return r, nil
}

func TestRoundTripper(t *testing.T) {
	downloadBody := ioutil.NopCloser(bytes.NewReader([]byte("data1")))
	uploadBody := ioutil.NopCloser(bytes.NewReader([]byte("data1")))

	base := &baseRoundTripper{
		responses: make(map[*http.Request]*http.Response),
	}
	downloadPool := &fakePool{}
	uploadPool := &fakePool{}
	rt := NewRoundTripper(base, downloadPool, uploadPool)

	// Empty request (no request, no reponse)
	uploadPool.reset()
	downloadPool.reset()
	req1, resp1 := base.add(&http.Request{}, &http.Response{})
	resp, err := rt.RoundTrip(req1)
	if resp != resp1 || err != nil {
		t.Errorf("invalid response or error: %v", err)
	}
	if len(downloadPool.readers) != 0 || len(uploadPool.readers) != 0 {
		t.Errorf("invalid pool contents: %v %v", downloadPool.readers, uploadPool.readers)
	}

	// Upload request
	uploadPool.reset()
	downloadPool.reset()
	req2, resp2 := base.add(&http.Request{
		Body: uploadBody,
	}, &http.Response{})
	resp, err = rt.RoundTrip(req2)
	if resp != resp2 || err != nil {
		t.Errorf("invalid response or error: %v", err)
	}
	if len(downloadPool.readers) != 0 || len(uploadPool.readers) != 1 {
		t.Errorf("invalid pool contents: %v %v", downloadPool.readers, uploadPool.readers)
	}

	// Download request
	uploadPool.reset()
	downloadPool.reset()
	req3, resp3 := base.add(&http.Request{}, &http.Response{Body: downloadBody})
	resp, err = rt.RoundTrip(req3)
	if resp != resp3 || err != nil {
		t.Errorf("invalid response or error: %v", err)
	}
	if len(downloadPool.readers) != 1 || len(uploadPool.readers) != 0 {
		t.Errorf("invalid pool contents: %v %v", downloadPool.readers, uploadPool.readers)
	}

	// Upload/Download request
	uploadPool.reset()
	downloadPool.reset()
	req4, resp4 := base.add(&http.Request{Body: uploadBody}, &http.Response{Body: downloadBody})
	resp, err = rt.RoundTrip(req4)
	if resp != resp4 || err != nil {
		t.Errorf("invalid response or error: %v", err)
	}
	if len(downloadPool.readers) != 1 || len(uploadPool.readers) != 1 {
		t.Errorf("invalid pool contents: %v %v", downloadPool.readers, uploadPool.readers)
	}
}