Skip to content

Commit

Permalink
Call updateKeyspaceMetadata in tests
Browse files Browse the repository at this point in the history 
This way more code is covered by tests. The tests now use
stretch/testify as dependency, which should be OK as it
was an indirect dependency before.
  • Loading branch information
@martin-sucha
martin-sucha committed Jul 2, 2019
1 parent fd1a383 commit f3de140
Show file tree
Hide file tree
Showing 4 changed files with 131 additions and 59 deletions.
7 changes: 7 additions & 0 deletions common_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"fmt"
"log"
"net"
"reflect"
"strings"
"sync"
"testing"
Expand Down Expand Up @@ -221,6 +222,12 @@ func assertEqual(t *testing.T, description string, expected, actual interface{})
}
}

func assertDeepEqual(t *testing.T, description string, expected, actual interface{}) {
if !reflect.DeepEqual(expected, actual) {
t.Errorf("expected %s to be (%+v) but was (%+v) instead", description, expected, actual)
}
}

func assertNil(t *testing.T, description string, actual interface{}) {
if actual != nil {
t.Errorf("expected %s to be (nil) but was (%+v) instead", description, actual)
Expand Down
16 changes: 7 additions & 9 deletions policies.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -427,7 +427,8 @@ type tokenAwareHostPolicy struct {
mu sync.RWMutex
partitioner string
fallback HostSelectionPolicy
session *Session
getKeyspaceMetadata func(keyspace string) (*KeyspaceMetadata, error)
getKeyspaceName func() string

tokenRing atomic.Value // *tokenRing
keyspaces atomic.Value // *keyspaceMeta
Expand All @@ -437,7 +438,8 @@ type tokenAwareHostPolicy struct {
}

func (t *tokenAwareHostPolicy) Init(s *Session) {
t.session = s
t.getKeyspaceMetadata = s.KeyspaceMetadata
t.getKeyspaceName = func() string {return s.cfg.Keyspace}
}

func (t *tokenAwareHostPolicy) IsLocal(host *HostInfo) bool {
Expand All @@ -459,7 +461,7 @@ func (t *tokenAwareHostPolicy) updateKeyspaceMetadata(keyspace string) {
replicas: make(map[string]map[token][]*HostInfo, size),
}

ks, err := t.session.KeyspaceMetadata(keyspace)
ks, err := t.getKeyspaceMetadata(keyspace)
if err == nil {
strat := getStrategy(ks)
if strat != nil {
Expand Down Expand Up @@ -495,16 +497,12 @@ func (t *tokenAwareHostPolicy) SetPartitioner(partitioner string) {

func (t *tokenAwareHostPolicy) AddHost(host *HostInfo) {
t.HostUp(host)
if t.session != nil { // disable for unit tests
t.updateKeyspaceMetadata(t.session.cfg.Keyspace)
}
t.updateKeyspaceMetadata(t.getKeyspaceName())
}

func (t *tokenAwareHostPolicy) RemoveHost(host *HostInfo) {
t.HostDown(host)
if t.session != nil { // disable for unit tests
t.updateKeyspaceMetadata(t.session.cfg.Keyspace)
}
t.updateKeyspaceMetadata(t.getKeyspaceName())
}

func (t *tokenAwareHostPolicy) HostUp(host *HostInfo) {
Expand Down
161 changes: 111 additions & 50 deletions policies_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
package gocql

import (
"errors"
"fmt"
"net"
"testing"
Expand Down Expand Up @@ -56,8 +57,12 @@ func TestHostPolicy_RoundRobin(t *testing.T) {
func TestHostPolicy_TokenAware(t *testing.T) {
policy := TokenAwareHostPolicy(RoundRobinHostPolicy())
policyInternal := policy.(*tokenAwareHostPolicy)

policyInternal.getKeyspaceName = func() string {return "myKeyspace"}
policyInternal.getKeyspaceMetadata = func(ks string) (*KeyspaceMetadata, error) {
return nil, errors.New("not initalized")
}
query := &Query{}
query.getKeyspace = func() string{return "myKeyspace"}

iter := policy.Pick(nil)
if iter == nil {
Expand Down Expand Up @@ -100,19 +105,31 @@ func TestHostPolicy_TokenAware(t *testing.T) {

policy.SetPartitioner("OrderedPartitioner")

// We need to simulate updateKeyspaceMetadata for replicas() to work.
// This corresponds to what simpleStrategy with rf=2 would generate.
newMeta := &keyspaceMeta{
replicas: map[string]map[token][]*HostInfo{
"": {
orderedToken("00"): {hosts[0], hosts[1]},
orderedToken("25"): {hosts[1], hosts[2]},
orderedToken("50"): {hosts[2], hosts[3]},
orderedToken("75"): {hosts[3], hosts[0]},
policyInternal.getKeyspaceMetadata = func(keyspaceName string) (*KeyspaceMetadata, error) {
if keyspaceName != "myKeyspace" {
return nil, fmt.Errorf("unknown keyspace: %s", keyspaceName)
}
return &KeyspaceMetadata{
Name: "myKeyspace",
StrategyClass: "SimpleStrategy",
StrategyOptions: map[string]interface{} {
"class": "SimpleStrategy",
"replication_factor": 2,
},
}, nil
}
policy.KeyspaceChanged(KeyspaceUpdateEvent{Keyspace: "myKeyspace"})

// The SimpleStrategy above should generate the following replicas.
// It's handy to have as reference here.
assertDeepEqual(t, "replicas", map[string]map[token][]*HostInfo{
"myKeyspace": {
orderedToken("00"): {hosts[0], hosts[1]},
orderedToken("25"): {hosts[1], hosts[2]},
orderedToken("50"): {hosts[2], hosts[3]},
orderedToken("75"): {hosts[3], hosts[0]},
},
}
policyInternal.keyspaces.Store(newMeta)
}, policyInternal.keyspaces.Load().(*keyspaceMeta).replicas)

// now the token ring is configured
query.RoutingKey([]byte("20"))
Expand Down Expand Up @@ -204,6 +221,11 @@ func TestHostPolicy_RoundRobin_NilHostInfo(t *testing.T) {

func TestHostPolicy_TokenAware_NilHostInfo(t *testing.T) {
policy := TokenAwareHostPolicy(RoundRobinHostPolicy())
policyInternal := policy.(*tokenAwareHostPolicy)
policyInternal.getKeyspaceName = func() string {return "myKeyspace"}
policyInternal.getKeyspaceMetadata = func(ks string) (*KeyspaceMetadata, error) {
return nil, errors.New("not initialized")
}

hosts := [...]*HostInfo{
{connectAddress: net.IPv4(10, 0, 0, 0), tokens: []string{"00"}},
Expand All @@ -217,6 +239,7 @@ func TestHostPolicy_TokenAware_NilHostInfo(t *testing.T) {
policy.SetPartitioner("OrderedPartitioner")

query := &Query{}
query.getKeyspace = func() string {return "myKeyspace"}
query.RoutingKey([]byte("20"))

iter := policy.Pick(query)
Expand Down Expand Up @@ -444,8 +467,13 @@ func TestHostPolicy_DCAwareRR(t *testing.T) {
func TestHostPolicy_TokenAware_DCAwareRR(t *testing.T) {
policy := TokenAwareHostPolicy(DCAwareRoundRobinPolicy("local"))
policyInternal := policy.(*tokenAwareHostPolicy)
policyInternal.getKeyspaceName = func() string {return "myKeyspace"}
policyInternal.getKeyspaceMetadata = func(ks string) (*KeyspaceMetadata, error) {
return nil, errors.New("not initialized")
}

query := &Query{}
query.getKeyspace = func() string {return "myKeyspace"}

iter := policy.Pick(nil)
if iter == nil {
Expand Down Expand Up @@ -488,27 +516,41 @@ func TestHostPolicy_TokenAware_DCAwareRR(t *testing.T) {

policy.SetPartitioner("OrderedPartitioner")

// we need to simulate updateKeyspaceMetadata for replicas() to work.
// this should correspond to what networkTopologyStrategy with rf=1 for each DC would generate
newMeta := &keyspaceMeta{
replicas: map[string]map[token][]*HostInfo{
"": {
orderedToken("05"): {hosts[0], hosts[1], hosts[2]},
orderedToken("10"): {hosts[1], hosts[2], hosts[3]},
orderedToken("15"): {hosts[2], hosts[3], hosts[4]},
orderedToken("20"): {hosts[3], hosts[4], hosts[5]},
orderedToken("25"): {hosts[4], hosts[5], hosts[6]},
orderedToken("30"): {hosts[5], hosts[6], hosts[7]},
orderedToken("35"): {hosts[6], hosts[7], hosts[8]},
orderedToken("40"): {hosts[7], hosts[8], hosts[9]},
orderedToken("45"): {hosts[8], hosts[9], hosts[10]},
orderedToken("50"): {hosts[9], hosts[10], hosts[11]},
orderedToken("55"): {hosts[10], hosts[11], hosts[0]},
orderedToken("60"): {hosts[11], hosts[0], hosts[1]},
policyInternal.getKeyspaceMetadata = func(keyspaceName string) (*KeyspaceMetadata, error) {
if keyspaceName != "myKeyspace" {
return nil, fmt.Errorf("unknown keyspace: %s", keyspaceName)
}
return &KeyspaceMetadata{
Name: "myKeyspace",
StrategyClass: "NetworkTopologyStrategy",
StrategyOptions: map[string]interface{} {
"class": "NetworkTopologyStrategy",
"local": 1,
"remote1": 1,
"remote2": 1,
},
}, nil
}
policy.KeyspaceChanged(KeyspaceUpdateEvent{Keyspace: "myKeyspace"})

// The NetworkTopologyStrategy above should generate the following replicas.
// It's handy to have as reference here.
assertDeepEqual(t, "replicas", map[string]map[token][]*HostInfo{
"myKeyspace": {
orderedToken("05"): {hosts[0], hosts[1], hosts[2]},
orderedToken("10"): {hosts[1], hosts[2], hosts[3]},
orderedToken("15"): {hosts[2], hosts[3], hosts[4]},
orderedToken("20"): {hosts[3], hosts[4], hosts[5]},
orderedToken("25"): {hosts[4], hosts[5], hosts[6]},
orderedToken("30"): {hosts[5], hosts[6], hosts[7]},
orderedToken("35"): {hosts[6], hosts[7], hosts[8]},
orderedToken("40"): {hosts[7], hosts[8], hosts[9]},
orderedToken("45"): {hosts[8], hosts[9], hosts[10]},
orderedToken("50"): {hosts[9], hosts[10], hosts[11]},
orderedToken("55"): {hosts[10], hosts[11], hosts[0]},
orderedToken("60"): {hosts[11], hosts[0], hosts[1]},
},
}
policyInternal.keyspaces.Store(newMeta)
}, policyInternal.keyspaces.Load().(*keyspaceMeta).replicas)

// now the token ring is configured
query.RoutingKey([]byte("23"))
Expand Down Expand Up @@ -544,8 +586,13 @@ func TestHostPolicy_TokenAware_DCAwareRR(t *testing.T) {
func TestHostPolicy_TokenAware_DCAwareRR_NonLocalFallback(t *testing.T) {
policy := TokenAwareHostPolicy(DCAwareRoundRobinPolicy("local"), NonLocalReplicasFallback())
policyInternal := policy.(*tokenAwareHostPolicy)
policyInternal.getKeyspaceName = func() string {return "myKeyspace"}
policyInternal.getKeyspaceMetadata = func(ks string) (*KeyspaceMetadata, error) {
return nil, errors.New("not initialized")
}

query := &Query{}
query.getKeyspace = func() string {return "myKeyspace"}

iter := policy.Pick(nil)
if iter == nil {
Expand Down Expand Up @@ -588,27 +635,41 @@ func TestHostPolicy_TokenAware_DCAwareRR_NonLocalFallback(t *testing.T) {

policy.SetPartitioner("OrderedPartitioner")

// we need to simulate updateKeyspaceMetadata for replicas() to work.
// this should correspond to what networkTopologyStrategy with rf=1 for each DC would generate
newMeta := &keyspaceMeta{
replicas: map[string]map[token][]*HostInfo{
"": {
orderedToken("05"): {hosts[0], hosts[1], hosts[2]},
orderedToken("10"): {hosts[1], hosts[2], hosts[3]},
orderedToken("15"): {hosts[2], hosts[3], hosts[4]},
orderedToken("20"): {hosts[3], hosts[4], hosts[5]},
orderedToken("25"): {hosts[4], hosts[5], hosts[6]},
orderedToken("30"): {hosts[5], hosts[6], hosts[7]},
orderedToken("35"): {hosts[6], hosts[7], hosts[8]},
orderedToken("40"): {hosts[7], hosts[8], hosts[9]},
orderedToken("45"): {hosts[8], hosts[9], hosts[10]},
orderedToken("50"): {hosts[9], hosts[10], hosts[11]},
orderedToken("55"): {hosts[10], hosts[11], hosts[0]},
orderedToken("60"): {hosts[11], hosts[0], hosts[1]},
policyInternal.getKeyspaceMetadata = func(keyspaceName string) (*KeyspaceMetadata, error) {
if keyspaceName != "myKeyspace" {
return nil, fmt.Errorf("unknown keyspace: %s", keyspaceName)
}
return &KeyspaceMetadata{
Name: "myKeyspace",
StrategyClass: "NetworkTopologyStrategy",
StrategyOptions: map[string]interface{} {
"class": "NetworkTopologyStrategy",
"local": 1,
"remote1": 1,
"remote2": 1,
},
}, nil
}
policy.KeyspaceChanged(KeyspaceUpdateEvent{Keyspace: "myKeyspace"})

// The NetworkTopologyStrategy above should generate the following replicas.
// It's handy to have as reference here.
assertDeepEqual(t, "replicas", map[string]map[token][]*HostInfo{
"myKeyspace": {
orderedToken("05"): {hosts[0], hosts[1], hosts[2]},
orderedToken("10"): {hosts[1], hosts[2], hosts[3]},
orderedToken("15"): {hosts[2], hosts[3], hosts[4]},
orderedToken("20"): {hosts[3], hosts[4], hosts[5]},
orderedToken("25"): {hosts[4], hosts[5], hosts[6]},
orderedToken("30"): {hosts[5], hosts[6], hosts[7]},
orderedToken("35"): {hosts[6], hosts[7], hosts[8]},
orderedToken("40"): {hosts[7], hosts[8], hosts[9]},
orderedToken("45"): {hosts[8], hosts[9], hosts[10]},
orderedToken("50"): {hosts[9], hosts[10], hosts[11]},
orderedToken("55"): {hosts[10], hosts[11], hosts[0]},
orderedToken("60"): {hosts[11], hosts[0], hosts[1]},
},
}
policyInternal.keyspaces.Store(newMeta)
}, policyInternal.keyspaces.Load().(*keyspaceMeta).replicas)

// now the token ring is configured
query.RoutingKey([]byte("18"))
Expand Down
6 changes: 6 additions & 0 deletions session.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -689,6 +689,9 @@ type Query struct {
metrics *queryMetrics

disableAutoPage bool

// getKeyspace is field so that it can be overriden in tests
getKeyspace func() string
}

func (q *Query) defaultsFromSession() {
Expand Down Expand Up @@ -909,6 +912,9 @@ func (q *Query) retryPolicy() RetryPolicy {

// Keyspace returns the keyspace the query will be executed against.
func (q *Query) Keyspace() string {
if q.getKeyspace != nil {
return q.getKeyspace()
}
if q.session == nil {
return ""
}
Expand Down

0 comments on commit f3de140

Please sign in to comment.