diff --git a/pkg/ipfs/mocks/dag_putters.go b/pkg/ipfs/mocks/dag_putters.go
index 1e5b26ae..7461344c 100644
--- a/pkg/ipfs/mocks/dag_putters.go
+++ b/pkg/ipfs/mocks/dag_putters.go
@@ -16,7 +16,11 @@
 
 package mocks
 
-import "errors"
+import (
+	"errors"
+
+	"github.com/ethereum/go-ethereum/common"
+)
 
 // DagPutter is a mock for testing the ipfs publisher
 type DagPutter struct {
@@ -29,19 +33,21 @@ func (dp *DagPutter) DagPut(raw interface{}) ([]string, error) {
 	return dp.CIDsToReturn, dp.ErrToReturn
 }
 
-// IncrementingDagPutter is a mock for testing the ipfs publisher
-type IncrementingDagPutter struct {
-	CIDsToReturn []string
-	iterator     int
+// MappedDagPutter is a mock for testing the ipfs publisher
+type MappedDagPutter struct {
+	CIDsToReturn map[common.Hash][]string
 	ErrToReturn  error
 }
 
 // DagPut returns the pre-loaded CIDs or error
-func (dp *IncrementingDagPutter) DagPut(raw interface{}) ([]string, error) {
-	if len(dp.CIDsToReturn) >= dp.iterator+1 {
-		cid := dp.CIDsToReturn[dp.iterator]
-		dp.iterator++
-		return []string{cid}, dp.ErrToReturn
+func (mdp *MappedDagPutter) DagPut(raw interface{}) ([]string, error) {
+	if mdp.CIDsToReturn == nil {
+		return nil, errors.New("mapped dag putter needs to be initialized with a map of cids to return")
 	}
-	return nil, errors.New("dag putter iterator is out of range")
+	by, ok := raw.([]byte)
+	if !ok {
+		return nil, errors.New("mapped dag putters can only dag put []byte values")
+	}
+	hash := common.BytesToHash(by)
+	return mdp.CIDsToReturn[hash], nil
 }
diff --git a/pkg/ipfs/publisher_test.go b/pkg/ipfs/publisher_test.go
index 7d45d807..87f5164b 100644
--- a/pkg/ipfs/publisher_test.go
+++ b/pkg/ipfs/publisher_test.go
@@ -17,6 +17,7 @@
 package ipfs_test
 
 import (
+	"github.com/ethereum/go-ethereum/common"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
 
@@ -28,8 +29,8 @@ var (
 	mockHeaderDagPutter  *mocks.DagPutter
 	mockTrxDagPutter     *mocks.DagPutter
 	mockRctDagPutter     *mocks.DagPutter
-	mockStateDagPutter   *mocks.IncrementingDagPutter
-	mockStorageDagPutter *mocks.IncrementingDagPutter
+	mockStateDagPutter   *mocks.MappedDagPutter
+	mockStorageDagPutter *mocks.DagPutter
 )
 
 var _ = Describe("Publisher", func() {
@@ -37,8 +38,8 @@ var _ = Describe("Publisher", func() {
 		mockHeaderDagPutter = new(mocks.DagPutter)
 		mockTrxDagPutter = new(mocks.DagPutter)
 		mockRctDagPutter = new(mocks.DagPutter)
-		mockStateDagPutter = new(mocks.IncrementingDagPutter)
-		mockStorageDagPutter = new(mocks.IncrementingDagPutter)
+		mockStateDagPutter = new(mocks.MappedDagPutter)
+		mockStorageDagPutter = new(mocks.DagPutter)
 	})
 
 	Describe("Publish", func() {
@@ -46,7 +47,12 @@ var _ = Describe("Publisher", func() {
 			mockHeaderDagPutter.CIDsToReturn = []string{"mockHeaderCID"}
 			mockTrxDagPutter.CIDsToReturn = []string{"mockTrxCID1", "mockTrxCID2"}
 			mockRctDagPutter.CIDsToReturn = []string{"mockRctCID1", "mockRctCID2"}
-			mockStateDagPutter.CIDsToReturn = []string{"mockStateCID1", "mockStateCID2"}
+			val1 := common.BytesToHash(mocks.MockIPLDPayload.StateNodes[mocks.ContractLeafKey].Value)
+			val2 := common.BytesToHash(mocks.MockIPLDPayload.StateNodes[mocks.AnotherContractLeafKey].Value)
+			mockStateDagPutter.CIDsToReturn = map[common.Hash][]string{
+				val1: {"mockStateCID1"},
+				val2: {"mockStateCID2"},
+			}
 			mockStorageDagPutter.CIDsToReturn = []string{"mockStorageCID"}
 			publisher := ipfs.Publisher{
 				HeaderPutter:      mockHeaderDagPutter,