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(Limits{42 * 1024, 42 * 1024}) limbe := LimitBackend(be, limiter) rd := restic.NewByteReader(data, nil) 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(Limits{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) } }