From 159bd9278e147b2a402f3cfe0fe5397e5d436d5d Mon Sep 17 00:00:00 2001
From: Gustav Westling <hello@gustav.tv>
Date: Wed, 23 Nov 2016 14:34:56 +0100
Subject: [PATCH] Fix bug where helper types are not properly initialized by
 using Set(). The fix requires that Key() is used before Set(). Will properly
 throw a warning if Exec() on helper types does not have a key set

---
 auto_map.go |  1 +
 counter.go  |  5 ++++
 set.go      |  5 ++++
 set_test.go | 72 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 83 insertions(+)

diff --git a/auto_map.go b/auto_map.go
index f6a1d11..10a3e8e 100644
--- a/auto_map.go
+++ b/auto_map.go
@@ -53,6 +53,7 @@ func (c Command) Set(val interface{}) Command {
 	riakContext, op, err := encodeInterface(val, requestData{
 		bucket:     c.bucket,
 		bucketType: c.bucketType,
+		key:        c.key,
 	})
 
 	if err != nil {
diff --git a/counter.go b/counter.go
index 0abf4bf..2d20cb7 100644
--- a/counter.go
+++ b/counter.go
@@ -59,6 +59,11 @@ func (c *Counter) Exec(client *Session) error {
 		return errors.New("Unknown path to counter. Retrieve counter with GetMap before updating the counter")
 	}
 
+	// Validate c.key
+	if c.key.bucket == "" || c.key.bucketType == "" || c.key.key == "" {
+		return errors.New("Invalid key in Counter Exec()")
+	}
+
 	op := &riak.MapOperation{}
 	outerOp := op
 
diff --git a/set.go b/set.go
index f6aa569..b00a96f 100644
--- a/set.go
+++ b/set.go
@@ -147,6 +147,11 @@ func (s *Set) Exec(client *Session) error {
 		return errors.New("Unknown path to Set. Retrieve Set with GetMap before updating the Set")
 	}
 
+	// Validate s.key
+	if s.key.bucket == "" || s.key.bucketType == "" || s.key.key == "" {
+		return errors.New("Invalid key in Set Exec()")
+	}
+
 	op := &riak.MapOperation{}
 	outerOp := op
 
diff --git a/set_test.go b/set_test.go
index 6ab6f2d..5df7894 100644
--- a/set_test.go
+++ b/set_test.go
@@ -1,6 +1,7 @@
 package goriak
 
 import (
+	"math/rand"
 	"reflect"
 	"sort"
 	"testing"
@@ -346,3 +347,74 @@ func TestSetInitializeSet(t *testing.T) {
 		t.Error("testVal2: Foos is not initialized")
 	}
 }
+
+var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
+
+func randomKey() string {
+	n := 10
+	b := make([]rune, n)
+	for i := range b {
+		b[i] = letters[rand.Intn(len(letters))]
+	}
+	return string(b)
+}
+
+func TestSetUnitialized(t *testing.T) {
+	type testType struct {
+		Foos *Set
+	}
+
+	val := testType{}
+	key := randomKey()
+
+	_, err := bucket().Key(key).Set(&val).Run(con())
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	err = val.Foos.AddString("wohoohooo").Exec(con())
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	// Fetch
+	var val2 testType
+	_, err = bucket().Get(key, &val2).Run(con())
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	if !val2.Foos.HasString("wohoohooo") {
+		t.Error("Not in set after save")
+	}
+}
+
+func TestSetUnitializedOtherOrder(t *testing.T) {
+	type testType struct {
+		Foos *Set
+	}
+
+	val := testType{}
+	key := randomKey()
+
+	// The difference is here
+	_, err := bucket().Set(&val).Key(key).Run(con())
+
+	if err != nil {
+		t.Error(err)
+	}
+
+	err = val.Foos.AddString("wohoohooo").Exec(con())
+
+	if err == nil {
+		t.Error("Got no error after AddString()")
+		return
+	}
+
+	if err.Error() != "Invalid key in Set Exec()" {
+		t.Error("Unexpected error:", err.Error())
+	}
+}