restic/internal/limiter/limiter_backend_test.go

110 lines
2.8 KiB
Go

package limiter
import (
"bytes"
"context"
"crypto/rand"
"fmt"
"io"
"testing"
"github.com/restic/restic/internal/mock"
"github.com/restic/restic/internal/restic"
rtest "github.com/restic/restic/internal/test"
)
func randomBytes(t *testing.T, size int) []byte {
data := make([]byte, size)
_, err := io.ReadFull(rand.Reader, data)
rtest.OK(t, err)
return data
}
func TestLimitBackendSave(t *testing.T) {
testHandle := restic.Handle{Type: restic.PackFile, Name: "test"}
data := randomBytes(t, 1234)
be := mock.NewBackend()
be.SaveFn = func(ctx context.Context, h restic.Handle, rd restic.RewindReader) error {
buf := new(bytes.Buffer)
_, err := io.Copy(buf, rd)
if err != nil {
return nil
}
if !bytes.Equal(data, buf.Bytes()) {
return fmt.Errorf("data mismatch")
}
return nil
}
limiter := NewStaticLimiter(42*1024, 42*1024)
limbe := LimitBackend(be, limiter)
rd := restic.NewByteReader(data)
err := limbe.Save(context.TODO(), testHandle, rd)
rtest.OK(t, err)
}
type tracedReadWriteToCloser struct {
io.Reader
io.WriterTo
Traced bool
}
func newTracedReadWriteToCloser(rd *bytes.Reader) *tracedReadWriteToCloser {
return &tracedReadWriteToCloser{Reader: rd, WriterTo: rd}
}
func (r *tracedReadWriteToCloser) WriteTo(w io.Writer) (n int64, err error) {
r.Traced = true
return r.WriterTo.WriteTo(w)
}
func (r *tracedReadWriteToCloser) Close() error {
return nil
}
func TestLimitBackendLoad(t *testing.T) {
testHandle := restic.Handle{Type: restic.PackFile, Name: "test"}
data := randomBytes(t, 1234)
for _, test := range []struct {
innerWriteTo, outerWriteTo bool
}{{false, false}, {false, true}, {true, false}, {true, true}} {
be := mock.NewBackend()
src := newTracedReadWriteToCloser(bytes.NewReader(data))
be.OpenReaderFn = func(ctx context.Context, h restic.Handle, length int, offset int64) (io.ReadCloser, error) {
if length != 0 || offset != 0 {
return nil, fmt.Errorf("Not supported")
}
// test both code paths in WriteTo of limitedReadCloser
if test.innerWriteTo {
return src, nil
}
return newTracedReadCloser(src), nil
}
limiter := NewStaticLimiter(42*1024, 42*1024)
limbe := LimitBackend(be, limiter)
err := limbe.Load(context.TODO(), testHandle, 0, 0, func(rd io.Reader) error {
dataRead := new(bytes.Buffer)
// test both Read and WriteTo
if !test.outerWriteTo {
rd = newTracedReadCloser(rd)
}
_, err := io.Copy(dataRead, rd)
if err != nil {
return err
}
if !bytes.Equal(data, dataRead.Bytes()) {
return fmt.Errorf("read broken data")
}
return nil
})
rtest.OK(t, err)
rtest.Assert(t, src.Traced == (test.innerWriteTo && test.outerWriteTo),
"unexpected/missing writeTo call innerWriteTo %v outerWriteTo %v",
test.innerWriteTo, test.outerWriteTo)
}
}