1 files changed, 46 insertions, 21 deletions
diff --git a/state/state_test.go b/state/state_test.go
index 737815e90..825d21fcc 100644
--- a/state/state_test.go
+++ b/state/state_test.go
@@ -1,36 +1,61 @@
 package state
 
 import (
-	"testing"
+	checker "gopkg.in/check.v1"
 
 	"github.com/ethereum/go-ethereum/ethdb"
 	"github.com/ethereum/go-ethereum/ethutil"
 	"github.com/ethereum/go-ethereum/trie"
 )
 
-var ZeroHash256 = make([]byte, 32)
-
-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 _ = checker.Suite(&StateSuite{})
 
-	snapshot := state.Copy()
+// var ZeroHash256 = make([]byte, 32)
 
-	stateObject = state.GetStateObject([]byte("aa"))
-	stateObject.SetStorage(ethutil.Big("0"), ethutil.NewValue(43))
+func (s *StateSuite) TestDump(c *checker.C) {
+	key := []byte{0x01}
+	value := "foo"
+	node := []interface{}{key, value}
+	s.state.Trie.Put(node)
+	dump := s.state.Dump()
+	c.Assert(dump, checker.NotNil)
+}
 
-	state.Set(snapshot)
+func (s *StateSuite) SetUpTest(c *checker.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 *checker.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, checker.DeepEquals, res)
 }