fix keyspace integration test

Author: joao-r-reis

keyspace_table_test.go

@@ -33,7 +33,7 @@ import (
 	"testing"
 )
 
-// Keyspace_table checks if Query.Keyspace() is updated based on prepared statement
+// Keyspace_table checks if Iter.Keyspace() is updated based on prepared statement
 func TestKeyspaceTable(t *testing.T) {
 	cluster := createCluster()
 
@@ -45,8 +45,7 @@ func TestKeyspaceTable(t *testing.T) {
 		t.Fatal("createSession:", err)
 	}
 
-	cluster.Keyspace = "wrong_keyspace"
-
+	wrongKeyspace := "testwrong"
 	keyspace := "test1"
 	table := "table1"
 
@@ -55,6 +54,11 @@ func TestKeyspaceTable(t *testing.T) {
 		t.Fatal("unable to drop keyspace:", err)
 	}
 
+	err = createTable(session, `DROP KEYSPACE IF EXISTS `+wrongKeyspace)
+	if err != nil {
+		t.Fatal("unable to drop keyspace:", err)
+	}
+
 	err = createTable(session, fmt.Sprintf(`CREATE KEYSPACE %s
 	WITH replication = {
 		'class' : 'SimpleStrategy',
@@ -65,6 +69,16 @@ func TestKeyspaceTable(t *testing.T) {
 		t.Fatal("unable to create keyspace:", err)
 	}
 
+	err = createTable(session, fmt.Sprintf(`CREATE KEYSPACE %s
+	WITH replication = {
+		'class' : 'SimpleStrategy',
+		'replication_factor' : 1
+	}`, wrongKeyspace))
+
+	if err != nil {
+		t.Fatal("unable to create keyspace:", err)
+	}
+
 	if err := session.control.awaitSchemaAgreement(); err != nil {
 		t.Fatal(err)
 	}
@@ -80,11 +94,24 @@ func TestKeyspaceTable(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	session.Close()
+
+	cluster = createCluster()
+
+	fallback = RoundRobinHostPolicy()
+	cluster.PoolConfig.HostSelectionPolicy = TokenAwareHostPolicy(fallback)
+	cluster.Keyspace = wrongKeyspace
+
+	session, err = cluster.CreateSession()
+	if err != nil {
+		t.Fatal("createSession:", err)
+	}
+
 	ctx := context.Background()
 
 	// insert a row
 	if err := session.Query(`INSERT INTO test1.table1(pk, ck, v) VALUES (?, ?, ?)`,
-		1, 2, 3).WithContext(ctx).Consistency(One).Exec(); err != nil {
+		1, 2, 3).Consistency(One).ExecContext(ctx); err != nil {
 		t.Fatal(err)
 	}
 
@@ -93,8 +120,8 @@ func TestKeyspaceTable(t *testing.T) {
 	/* Search for a specific set of records whose 'pk' column matches
 	 * the value of inserted row. */
 	qry := session.Query(`SELECT pk FROM test1.table1 WHERE pk = ? LIMIT 1`,
-		1).WithContext(ctx).Consistency(One)
-	iter := qry.Iter()
+		1).Consistency(One)
+	iter := qry.IterContext(ctx)
 	ok := iter.Scan(&pk)
 	err = iter.Close()
 	if err != nil {
@@ -104,9 +131,9 @@ func TestKeyspaceTable(t *testing.T) {
 		t.Fatal("expected pk to be scanned")
 	}
 
-	// cluster.Keyspace was set to "wrong_keyspace", but during prepering statement
+	// cluster.Keyspace was set to "testwrong", but during prepering statement
 	// Keyspace in Query should be changed to "test" and Table should be changed to table1
-	assertEqual(t, "qry.Keyspace()", "test1", qry.Keyspace())
+	assertEqual(t, "qry.Keyspace()", "testwrong", qry.Keyspace())
 	assertEqual(t, "iter.Keyspace()", "test1", iter.Keyspace())
 	assertEqual(t, "iter.Table()", "table1", iter.Table())
 }