1 files changed, 40 insertions, 19 deletions
diff --git a/state/state_test.go b/state/state_test.go
index 737815e90..9dd965e95 100644
--- a/state/state_test.go
+++ b/state/state_test.go
@@ -1,6 +1,7 @@
 package state
 
 import (
+	. "gopkg.in/check.v1"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/ethdb"
@@ -8,29 +9,49 @@ import (
 	"github.com/ethereum/go-ethereum/trie"
 )
 
-var ZeroHash256 = make([]byte, 32)
+func Test(t *testing.T) { TestingT(t) }
 
-func TestSnapshot(t *testing.T) {
-	db, _ := ethdb.NewMemDatabase()
-	ethutil.ReadConfig(".ethtest", "/tmp/ethtest", "")
-	ethutil.Config.Db = db
-
-	state := New(trie.New(db, ""))
-
-	stateObject := state.GetOrNewStateObject([]byte("aa"))
+type StateSuite struct {
+	state *State
+}
 
-	stateObject.SetStorage(ethutil.Big("0"), ethutil.NewValue(42))
+var _ = Suite(&StateSuite{})
 
-	snapshot := state.Copy()
+const expectedasbytes = "Expected % x Got % x"
 
-	stateObject = state.GetStateObject([]byte("aa"))
-	stateObject.SetStorage(ethutil.Big("0"), ethutil.NewValue(43))
+// var ZeroHash256 = make([]byte, 32)
 
-	state.Set(snapshot)
+func (s *StateSuite) SetUpTest(c *C) {
+	db, _ := ethdb.NewMemDatabase()
+	ethutil.ReadConfig(".ethtest", "/tmp/ethtest", "")
+	ethutil.Config.Db = db
+	s.state = New(trie.New(db, ""))
+}
 
-	stateObject = state.GetStateObject([]byte("aa"))
-	res := stateObject.GetStorage(ethutil.Big("0"))
-	if !res.Cmp(ethutil.NewValue(42)) {
-		t.Error("Expected storage 0 to be 42", res)
-	}
+func (s *StateSuite) TestSnapshot(c *C) {
+	stateobjaddr := []byte("aa")
+	storageaddr := ethutil.Big("0")
+	data1 := ethutil.NewValue(42)
+	data2 := ethutil.NewValue(43)
+
+	// get state object
+	stateObject := s.state.GetOrNewStateObject(stateobjaddr)
+	// set inital state object value
+	stateObject.SetStorage(storageaddr, data1)
+	// get snapshot of current state
+	snapshot := s.state.Copy()
+
+	// get state object. is this strictly necessary?
+	stateObject = s.state.GetStateObject(stateobjaddr)
+	// set new state object value
+	stateObject.SetStorage(storageaddr, data2)
+	// restore snapshot
+	s.state.Set(snapshot)
+
+	// get state object
+	stateObject = s.state.GetStateObject(stateobjaddr)
+	// get state storage value
+	res := stateObject.GetStorage(storageaddr)
+
+	c.Assert(data1, DeepEquals, res, Commentf(expectedasbytes, data1, res))
 }