forked from pachyderm/pachyderm
-
Notifications
You must be signed in to change notification settings - Fork 0
/
mock_client.go
120 lines (107 loc) · 2.36 KB
/
mock_client.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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
package discovery
import (
"path"
"strings"
"sync"
"github.com/pachyderm/pachyderm/src/pfs"
)
type record struct {
directory bool
data string
}
type mockClient struct {
records map[string]record
lock sync.RWMutex
}
func newMockClient() *mockClient {
return &mockClient{
make(map[string]record),
sync.RWMutex{},
}
}
func (c *mockClient) Close() error {
return nil
}
func (c *mockClient) Get(key string) (string, error) {
c.lock.RLock()
defer c.lock.RUnlock()
record, ok := c.records[key]
if !ok {
return "", pfs.ErrDiscoveryNotFound
}
if record.directory {
return "", pfs.ErrDiscoveryNotValue
}
return record.data, nil
}
func (c *mockClient) GetAll(key string) (map[string]string, error) {
c.lock.RLock()
defer c.lock.RUnlock()
result := make(map[string]string)
for k, v := range c.records {
if strings.HasPrefix(k, key) && !v.directory {
result[k] = v.data
}
}
return result, nil
}
func (c *mockClient) Set(key string, value string) error {
c.lock.Lock()
defer c.lock.Unlock()
return c.unsafeSet(key, value)
}
func (c *mockClient) Create(key string, value string) error {
c.lock.Lock()
defer c.lock.Unlock()
_, ok := c.records[key]
if ok {
return pfs.ErrDiscoveryKeyAlreadyExists
}
return c.unsafeSet(key, value)
}
func (c *mockClient) Delete(key string) error {
c.lock.Lock()
defer c.lock.Unlock()
oldRecord, ok := c.records[key]
if !ok {
return nil
}
if oldRecord.directory {
return pfs.ErrDiscoveryNotValue
}
delete(c.records, key)
return nil
}
func (c *mockClient) CheckAndSet(key string, value string, oldValue string) error {
c.lock.Lock()
defer c.lock.Unlock()
oldRecord, ok := c.records[key]
if !ok {
return pfs.ErrDiscoveryNotFound
}
if oldRecord.directory {
return pfs.ErrDiscoveryNotValue
}
if oldRecord.data != oldValue {
return pfs.ErrDiscoveryPreconditionNotMet
}
return c.unsafeSet(key, value)
}
func (c *mockClient) unsafeSet(key string, value string) error {
parts := strings.Split(key, "/")
for i, _ := range parts[:len(parts)-1] {
oldRecord, ok := c.records["/"+path.Join(parts[:i]...)]
if ok && oldRecord.directory {
return pfs.ErrDiscoveryNotDirectory
}
if !ok {
c.records["/"+path.Join(parts[:i]...)] = record{true, ""}
}
}
oldRecord, ok := c.records[key]
if ok && oldRecord.directory {
return pfs.ErrDiscoveryNotValue
}
c.records[key] = record{false, value}
return nil
}