wal-g

1
package testutils
2

3
import (
4
	"io"
5
	"sync"
6

7
	"github.com/wal-g/wal-g/utility"
8

9
	"github.com/wal-g/wal-g/pkg/storages/storage"
10
)
11

12
func NewUnexpectedEOFLimitReader(readCloser io.ReadCloser, maxRead int64) io.ReadCloser {
13
	return &UnexpectedEOFLimitReader{
14
		source: readCloser,
15
		left:   maxRead,
16
	}
17
}
18

19
type UnexpectedEOFLimitReader struct {
20
	source io.ReadCloser
21
	left   int64
22
}
23

24
func (u *UnexpectedEOFLimitReader) Read(p []byte) (n int, err error) {
25
	if u.left <= 0 {
26
		return 0, io.ErrUnexpectedEOF
27
	}
28
	unexpectedEOF := false
29
	if int64(len(p)) > u.left {
30
		unexpectedEOF = true
31
		p = p[:u.left]
32
	}
33
	n, err = u.source.Read(p)
34
	u.left -= int64(n)
35
	if unexpectedEOF && err == nil {
36
		err = io.ErrUnexpectedEOF
37
	}
38
	return
39
}
40

41
func (u *UnexpectedEOFLimitReader) Close() error {
42
	return u.source.Close()
43
}
44

45
func NewNetworkErrorFolder(sourceFolder storage.Folder, networkErrorAfterByteSize int) storage.Folder {
46
	return &TestFolder{
47
		Folder:       sourceFolder,
48
		maxReadSize:  int64(networkErrorAfterByteSize),
49
		readFromFile: make(map[string]*int64),
50
	}
51
}
52

53
type TestFolder struct {
54
	storage.Folder
55
	maxReadSize  int64
56
	readFromFile map[string]*int64
57
	mutex        sync.Mutex
58
}
59

60
func (tf *TestFolder) ReadObject(path string) (io.ReadCloser, error) {
61
	reader, err := tf.Folder.ReadObject(path)
62
	if tf.maxReadSize != 0 {
63
		tf.mutex.Lock()
64
		defer tf.mutex.Unlock()
65
		if _, ok := tf.readFromFile[path]; !ok {
66
			tf.readFromFile[path] = new(int64)
67
		}
68
		reader = utility.NewWithSizeReadCloser(reader, tf.readFromFile[path])
69
		reader = NewUnexpectedEOFLimitReader(reader, *tf.readFromFile[path]+tf.maxReadSize)
70
	}
71
	return reader, err
72
}
73