Add CidIndex to test framework

tests/integration/state.go

@@ -175,6 +175,9 @@ type state struct {
 	// nodes.
 	docIDs [][]client.DocID
 
+	// CIDs by index, by field name, by doc index, by collection index.
+	cids [][]map[string][]cid.Cid
+
 	// Indexes, by index, by collection index, by node index.
 	indexes [][][]client.IndexDescription
 

tests/integration/test_case.go

@@ -17,6 +17,7 @@ import (
 	"github.com/sourcenetwork/immutable"
 
 	"github.com/sourcenetwork/defradb/client"
+	"github.com/sourcenetwork/defradb/internal/core"
 	"github.com/sourcenetwork/defradb/net"
 	"github.com/sourcenetwork/defradb/tests/gen"
 	"github.com/sourcenetwork/defradb/tests/predefined"
@@ -348,6 +349,50 @@ func NewDocIndex(collectionIndex int, index int) DocIndex {
 	}
 }
 
+// CidIndex represents a Cid previously stored in the node, it allows Cids to be referenced without worrying
+// about the specific bytes.
+type CidIndex struct {
+	// CollectionIndex is the index of the collection used to create the document that the Cid belongs to.
+	CollectionIndex int
+
+	// DocIndex is the index of the document that the Cid belongs to.
+	DocIndex int
+
+	// FieldName is the name of the field that this cid is for.
+	//
+	// Composite commits can be targetted with "C".
+	FieldName string
+
+	// The index of the Cid given the other parameters.
+	//
+	// This is typically equal to commit height - 1.
+	Index int
+}
+
+// NewCompCidIndex creates a new [CidIndex] instance targetting the cid of a composite commit.
+//
+// It allows cids to be referenced without worrying about the specific bytes that form it.
+func NewCompCidIndex(collectionIndex int, docIndex int, index int) CidIndex {
+	return CidIndex{
+		CollectionIndex: collectionIndex,
+		DocIndex:        docIndex,
+		FieldName:       core.COMPOSITE_NAMESPACE,
+		Index:           index,
+	}
+}
+
+// NewCidIndex creates a new [CidIndex] instance targetting a cid of a commit.
+//
+// It allows cids to be referenced without worrying about the specific bytes that form it.
+func NewCidIndex(collectionIndex int, docIndex int, fieldName string, index int) CidIndex {
+	return CidIndex{
+		CollectionIndex: collectionIndex,
+		DocIndex:        docIndex,
+		FieldName:       fieldName,
+		Index:           index,
+	}
+}
+
 // DeleteDoc will attempt to delete the given document in the given collection
 // using the collection api.
 type DeleteDoc struct {

tests/integration/utils.go

@@ -24,6 +24,8 @@ import (
 	"time"
 
 	"github.com/fxamacker/cbor/v2"
+	"github.com/ipfs/go-cid"
+	"github.com/ipfs/go-datastore/query"
 	"github.com/sourcenetwork/corelog"
 	"github.com/sourcenetwork/immutable"
 	"github.com/stretchr/testify/assert"
@@ -34,6 +36,7 @@ import (
 	"github.com/sourcenetwork/defradb/crypto"
 	"github.com/sourcenetwork/defradb/datastore"
 	"github.com/sourcenetwork/defradb/errors"
+	"github.com/sourcenetwork/defradb/internal/core"
 	"github.com/sourcenetwork/defradb/internal/db"
 	"github.com/sourcenetwork/defradb/internal/encryption"
 	"github.com/sourcenetwork/defradb/internal/request/graphql"
@@ -1271,8 +1274,59 @@ func createDoc(
 		// Expand the slice if required, so that the document can be accessed by collection index
 		s.docIDs = append(s.docIDs, make([][]client.DocID, action.CollectionID-len(s.docIDs)+1)...)
 	}
+	startDocIndex := len(s.docIDs[action.CollectionID])
 	s.docIDs[action.CollectionID] = append(s.docIDs[action.CollectionID], docIDs...)
 
+	if action.CollectionID >= len(s.cids) {
+		// Expand the slice if required, so that the cids can be accessed by collection index
+		s.cids = append(s.cids, make([][]map[string][]cid.Cid, action.CollectionID+1)...)
+	}
+
+	endDocIndex := startDocIndex + len(docIDs) - 1
+	if endDocIndex >= len(s.cids[action.CollectionID]) {
+		s.cids[action.CollectionID] = append(
+			s.cids[action.CollectionID],
+			make([]map[string][]cid.Cid, endDocIndex-len(s.cids[action.CollectionID])+1)...,
+		)
+	}
+
+	for relativeDocIndex, docID := range docIDs {
+		results, err := nodes[0].Headstore().Query(
+			s.ctx,
+			query.Query{
+				Prefix: core.HeadStoreKey{
+					DocID: docID.String(),
+				}.ToString(),
+				KeysOnly: true,
+			},
+		)
+		if err != nil {
+			s.t.Fatalf("failed to fetch cids: %s", err.Error())
+		}
+
+		for res := range results.Next() {
+			if res.Error != nil {
+				s.t.Fatalf("failed to fetch cids: %s", res.Error.Error())
+			}
+
+			key, err := core.NewHeadStoreKey(res.Key)
+			if err != nil {
+				s.t.Fatalf("failed to fetch cids: %s", err.Error())
+			}
+
+			docIndex := startDocIndex + relativeDocIndex
+
+			if s.cids[action.CollectionID][docIndex] == nil {
+				s.cids[action.CollectionID][docIndex] = map[string][]cid.Cid{}
+			}
+
+			// This is a new doc, so we can blindly overwrite the array of cids
+			s.cids[action.CollectionID][docIndex][key.FieldID] = []cid.Cid{
+				key.Cid,
+			}
+		}
+	}
+
 	if action.ExpectedError == "" {
 		waitForUpdateEvents(s, action.NodeID, getEventsForCreateDoc(s, action))
 	}
@@ -1453,6 +1507,35 @@ func deleteDoc(
 
 	assertExpectedErrorRaised(s.t, s.testCase.Description, action.ExpectedError, expectedErrorRaised)
 
+	results, err := nodes[0].Headstore().Query(
+		s.ctx,
+		query.Query{
+			Prefix: core.HeadStoreKey{
+				DocID: s.docIDs[action.CollectionID][action.DocID].String(),
+			}.ToString(),
+			KeysOnly: true,
+		},
+	)
+	if err != nil {
+		s.t.Fatalf("failed to fetch cids: %s", err.Error())
+	}
+
+	for res := range results.Next() {
+		if res.Error != nil {
+			s.t.Fatalf("failed to fetch cids: %s", res.Error.Error())
+		}
+
+		key, err := core.NewHeadStoreKey(res.Key)
+		if err != nil {
+			s.t.Fatalf("failed to fetch cids: %s", err.Error())
+		}
+
+		s.cids[action.CollectionID][action.DocID][key.FieldID] = append(
+			s.cids[action.CollectionID][action.DocID][key.FieldID],
+			key.Cid,
+		)
+	}
+
 	if action.ExpectedError == "" {
 		docIDs := map[string]struct{}{
 			docID.String(): {},
@@ -1504,6 +1587,35 @@ func updateDoc(
 	if action.ExpectedError == "" && !action.SkipLocalUpdateEvent {
 		waitForUpdateEvents(s, action.NodeID, getEventsForUpdateDoc(s, action))
 	}
+
+	results, err := nodes[0].Headstore().Query(
+		s.ctx,
+		query.Query{
+			Prefix: core.HeadStoreKey{
+				DocID: s.docIDs[action.CollectionID][action.DocID].String(),
+			}.ToString(),
+			KeysOnly: true,
+		},
+	)
+	if err != nil {
+		s.t.Fatalf("failed to fetch cids: %s", err.Error())
+	}
+
+	for res := range results.Next() {
+		if res.Error != nil {
+			s.t.Fatalf("failed to fetch cids: %s", res.Error.Error())
+		}
+
+		key, err := core.NewHeadStoreKey(res.Key)
+		if err != nil {
+			s.t.Fatalf("failed to fetch cids: %s", err.Error())
+		}
+
+		s.cids[action.CollectionID][action.DocID][key.FieldID] = append(
+			s.cids[action.CollectionID][action.DocID][key.FieldID],
+			key.Cid,
+		)
+	}
 }
 
 func updateDocViaColSave(
@@ -2079,6 +2191,18 @@ func assertRequestResultDocs(
 					actualValue,
 					fmt.Sprintf("node: %v, path: %s", nodeID, stack),
 				)
+
+			case CidIndex:
+				//nolint:lll
+				expectedCid := s.cids[expectedValue.CollectionIndex][expectedValue.DocIndex][expectedValue.FieldName][expectedValue.Index].String()
+				assertResultsEqual(
+					s.t,
+					s.clientType,
+					expectedCid,
+					actualValue,
+					fmt.Sprintf("node: %v, path: %s", nodeID, stack),
+				)
+
 			case []map[string]any:
 				actualValueMap := ConvertToArrayOfMaps(s.t, actualValue)