graph/db: factor out helper code

Here we factor out some of the crud code in TestEdgePolicyMissingMaxHtcl
so that we can re-use it later on.

Author: ellemouton

graph/db/graph_test.go

@@ -3734,12 +3734,12 @@ func TestDisabledChannelIDs(t *testing.T) {
 	}
 }
 
-// TestEdgePolicyMissingMaxHtcl tests that if we find a ChannelEdgePolicy in
+// TestEdgePolicyMissingMaxHTLC tests that if we find a ChannelEdgePolicy in
 // the DB that indicates that it should support the htlc_maximum_value_msat
 // field, but it is not part of the opaque data, then we'll handle it as it is
 // unknown. It also checks that we are correctly able to overwrite it when we
 // receive the proper update.
-func TestEdgePolicyMissingMaxHtcl(t *testing.T) {
+func TestEdgePolicyMissingMaxHTLC(t *testing.T) {
 	t.Parallel()
 	ctx := context.Background()
 
@@ -3797,45 +3797,10 @@ func TestEdgePolicyMissingMaxHtcl(t *testing.T) {
 	require.ErrorIs(t, err, ErrEdgePolicyOptionalFieldNotFound)
 
 	// Put the stripped bytes in the DB.
-	err = kvdb.Update(boltStore.db, func(tx kvdb.RwTx) error {
-		edges := tx.ReadWriteBucket(edgeBucket)
-		if edges == nil {
-			return ErrEdgeNotFound
-		}
-
-		edgeIndex := edges.NestedReadWriteBucket(edgeIndexBucket)
-		if edgeIndex == nil {
-			return ErrEdgeNotFound
-		}
-
-		var edgeKey [33 + 8]byte
-		copy(edgeKey[:], from)
-		byteOrder.PutUint64(edgeKey[33:], edge1.ChannelID)
-
-		var scratch [8]byte
-		var indexKey [8 + 8]byte
-		copy(indexKey[:], scratch[:])
-		byteOrder.PutUint64(indexKey[8:], edge1.ChannelID)
-
-		updateIndex, err := edges.CreateBucketIfNotExists(
-			edgeUpdateIndexBucket,
-		)
-		if err != nil {
-			return err
-		}
-
-		if err := updateIndex.Put(indexKey[:], nil); err != nil {
-			return err
-		}
-
-		return edges.Put(edgeKey[:], stripped)
-	}, func() {})
-	require.NoError(t, err, "error writing db")
+	putSerializedPolicy(t, boltStore.db, from, chanID, stripped)
 
 	// And add the second, unmodified edge.
-	if err := graph.UpdateEdgePolicy(ctx, edge2); err != nil {
-		t.Fatalf("unable to update edge: %v", err)
-	}
+	require.NoError(t, graph.UpdateEdgePolicy(ctx, edge2))
 
 	// Attempt to fetch the edge and policies from the DB. Since the policy
 	// we added is invalid according to the new format, it should be as we
@@ -3870,6 +3835,38 @@ func TestEdgePolicyMissingMaxHtcl(t *testing.T) {
 	assertEdgeInfoEqual(t, dbEdgeInfo, edgeInfo)
 }
 
+// putSerializedPolicy is a helper function that writes a serialized
+// ChannelEdgePolicy to the edge bucket in the database.
+func putSerializedPolicy(t *testing.T, db kvdb.Backend, from []byte,
+	chanID uint64, b []byte) {
+
+	err := kvdb.Update(db, func(tx kvdb.RwTx) error {
+		edges := tx.ReadWriteBucket(edgeBucket)
+		require.NotNil(t, edges)
+
+		edgeIndex := edges.NestedReadWriteBucket(edgeIndexBucket)
+		require.NotNil(t, edgeIndex)
+
+		var edgeKey [33 + 8]byte
+		copy(edgeKey[:], from)
+		byteOrder.PutUint64(edgeKey[33:], chanID)
+
+		var scratch [8]byte
+		var indexKey [8 + 8]byte
+		copy(indexKey[:], scratch[:])
+		byteOrder.PutUint64(indexKey[8:], chanID)
+
+		updateIndex, err := edges.CreateBucketIfNotExists(
+			edgeUpdateIndexBucket,
+		)
+		require.NoError(t, err)
+		require.NoError(t, updateIndex.Put(indexKey[:], nil))
+
+		return edges.Put(edgeKey[:], b)
+	}, func() {})
+	require.NoError(t, err, "error writing db")
+}
+
 // assertNumZombies queries the provided ChannelGraph for NumZombies, and
 // asserts that the returned number is equal to expZombies.
 func assertNumZombies(t *testing.T, graph *ChannelGraph, expZombies uint64) {