This repository has been archived by the owner on Oct 9, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 62
/
storage.go
79 lines (65 loc) · 2.54 KB
/
storage.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
package mocks
import (
"context"
"fmt"
"io"
"strings"
"github.com/flyteorg/flytestdlib/storage"
"github.com/golang/protobuf/proto"
)
type NopCloser struct {
io.Reader
}
func (NopCloser) Close() error { return nil }
type TestDataStore struct {
HeadCb func(ctx context.Context, reference storage.DataReference) (storage.Metadata, error)
ReadProtobufCb func(ctx context.Context, reference storage.DataReference, msg proto.Message) error
WriteProtobufCb func(
ctx context.Context, reference storage.DataReference, opts storage.Options, msg proto.Message) error
ConstructReferenceCb func(
ctx context.Context, reference storage.DataReference, nestedKeys ...string) (storage.DataReference, error)
Store map[storage.DataReference][]byte
}
func (t *TestDataStore) Head(ctx context.Context, reference storage.DataReference) (storage.Metadata, error) {
return t.HeadCb(ctx, reference)
}
func (t *TestDataStore) ReadProtobuf(ctx context.Context, reference storage.DataReference, msg proto.Message) error {
return t.ReadProtobufCb(ctx, reference, msg)
}
func (t *TestDataStore) WriteProtobuf(
ctx context.Context, reference storage.DataReference, opts storage.Options, msg proto.Message) error {
return t.WriteProtobufCb(ctx, reference, opts, msg)
}
func (t *TestDataStore) GetBaseContainerFQN(ctx context.Context) storage.DataReference {
return "s3://bucket"
}
// Retrieves a byte array from the Blob store or an error
func (t *TestDataStore) ReadRaw(ctx context.Context, reference storage.DataReference) (io.ReadCloser, error) {
return NopCloser{}, nil
}
// Stores a raw byte array.
func (t *TestDataStore) WriteRaw(
ctx context.Context, reference storage.DataReference, size int64, opts storage.Options, raw io.Reader) error {
return nil
}
// Copies from source to destination.
func (t *TestDataStore) CopyRaw(ctx context.Context, source, destination storage.DataReference, opts storage.Options) error {
return nil
}
func (t *TestDataStore) ConstructReference(
ctx context.Context, reference storage.DataReference, nestedKeys ...string) (storage.DataReference, error) {
if t.ConstructReferenceCb != nil {
return t.ConstructReferenceCb(ctx, reference, nestedKeys...)
}
nestedPath := strings.Join(nestedKeys, "/")
return storage.DataReference(fmt.Sprintf("%s/%v", reference, nestedPath)), nil
}
func GetMockStorageClient() *storage.DataStore {
mockStorageClient := TestDataStore{
Store: make(map[storage.DataReference][]byte),
}
return &storage.DataStore{
ComposedProtobufStore: &mockStorageClient,
ReferenceConstructor: &mockStorageClient,
}
}