forked from dolthub/vitess
-
Notifications
You must be signed in to change notification settings - Fork 0
/
testserver.go
56 lines (47 loc) · 1.74 KB
/
testserver.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
package zktopo
import (
"fmt"
"testing"
"github.com/youtube/vitess/go/vt/topo"
"github.com/youtube/vitess/go/zk"
"github.com/youtube/vitess/go/zk/fakezk"
"golang.org/x/net/context"
"launchpad.net/gozk/zookeeper"
)
// TestServer is a proxy for a real implementation of topo.Server that
// provides hooks for testing.
type TestServer struct {
topo.Impl
localCells []string
HookLockSrvShardForAction func()
}
// newTestServer returns a new TestServer (with the required paths created)
func newTestServer(t *testing.T, cells []string) topo.Impl {
zconn := fakezk.NewConn()
// create the toplevel zk paths
if _, err := zk.CreateRecursive(zconn, "/zk/global/vt", "", 0, zookeeper.WorldACL(zookeeper.PERM_ALL)); err != nil {
t.Fatalf("cannot init ZooKeeper: %v", err)
}
for _, cell := range cells {
if _, err := zk.CreateRecursive(zconn, fmt.Sprintf("/zk/%v/vt", cell), "", 0, zookeeper.WorldACL(zookeeper.PERM_ALL)); err != nil {
t.Fatalf("cannot init ZooKeeper: %v", err)
}
}
return &TestServer{Impl: &Server{zconn}, localCells: cells}
}
// NewTestServer returns a new TestServer (with the required paths created)
func NewTestServer(t *testing.T, cells []string) topo.Server {
return topo.Server{Impl: newTestServer(t, cells)}
}
// GetKnownCells is part of topo.Server interface
func (s *TestServer) GetKnownCells(ctx context.Context) ([]string, error) {
return s.localCells, nil
}
// LockSrvShardForAction should override the function defined by the underlying
// topo.Server.
func (s *TestServer) LockSrvShardForAction(ctx context.Context, cell, keyspace, shard, contents string) (string, error) {
if s.HookLockSrvShardForAction != nil {
s.HookLockSrvShardForAction()
}
return s.Impl.LockSrvShardForAction(ctx, cell, keyspace, shard, contents)
}