Fix tests

d2layouts/d2sequence/layout.go

@@ -58,7 +58,7 @@ func Layout(ctx context.Context, g *d2graph.Graph, layout func(ctx context.Conte
 	layoutObjects, objectOrder := getLayoutObjects(g, objectsToRemove)
 	g.Objects = layoutObjects
 
-	if g.Root.Attributes.Shape.Value == d2target.ShapeSequenceDiagram {
+	if isRootSequenceDiagram(g) {
 		// don't need to run the layout engine if the root is a sequence diagram
 		g.Root.TopLeft = geo.NewPoint(0, 0)
 	} else if err := layout(ctx, g); err != nil {
@@ -66,10 +66,13 @@ func Layout(ctx context.Context, g *d2graph.Graph, layout func(ctx context.Conte
 	}
 
 	cleanup(g, sequenceDiagrams, objectOrder, edgeOrder)
-
 	return nil
 }
 
+func isRootSequenceDiagram(g *d2graph.Graph) bool {
+	return g.Root.Attributes.Shape.Value == d2target.ShapeSequenceDiagram
+}
+
 // layoutSequenceDiagram finds the edges inside the sequence diagram and performs the layout on the object descendants
 func layoutSequenceDiagram(g *d2graph.Graph, obj *d2graph.Object) *sequenceDiagram {
 	// find the edges that belong to this sequence diagram
@@ -119,7 +122,13 @@ func getLayoutObjects(g *d2graph.Graph, toRemove map[*d2graph.Object]struct{}) (
 // - adds the sequence diagram descendants back to the graph objects
 // - sorts edges and objects to their original graph order
 func cleanup(g *d2graph.Graph, sequenceDiagrams map[string]*sequenceDiagram, objectsOrder, edgesOrder map[string]int) {
-	for _, obj := range g.Objects {
+	var objects []*d2graph.Object
+	if isRootSequenceDiagram(g) {
+		objects = []*d2graph.Object{g.Root}
+	} else {
+		objects = g.Objects
+	}
+	for _, obj := range objects {
 		if _, exists := sequenceDiagrams[obj.AbsID()]; !exists {
 			continue
 		}

d2layouts/d2sequence/layout_test.go

@@ -36,6 +36,7 @@ func TestBasicSequenceDiagram(t *testing.T) {
 		{
 			Src:   n1,
 			Dst:   n2,
+			Index: 0,
 			Attributes: d2graph.Attributes{
 				Label: d2graph.Scalar{Value: "left to right"},
 			},
@@ -43,6 +44,7 @@ func TestBasicSequenceDiagram(t *testing.T) {
 		{
 			Src:   n2,
 			Dst:   n1,
+			Index: 0,
 			Attributes: d2graph.Attributes{
 				Label: d2graph.Scalar{Value: "right to left"},
 			},
@@ -50,10 +52,12 @@ func TestBasicSequenceDiagram(t *testing.T) {
 		{
 			Src:   n1,
 			Dst:   n2,
+			Index: 1,
 		},
 		{
 			Src:   n2,
 			Dst:   n1,
+			Index: 1,
 		},
 	}
 	nEdges := len(g.Edges)
@@ -189,15 +193,19 @@ func TestSpansSequenceDiagram(t *testing.T) {
 		{
 			Src:   a_t1,
 			Dst:   b_t1,
+			Index: 0,
 		}, {
 			Src:   b_t1,
 			Dst:   a_t1,
+			Index: 0,
 		}, {
 			Src:   a_t2,
 			Dst:   b,
+			Index: 0,
 		}, {
 			Src:   b,
 			Dst:   a_t2,
+			Index: 0,
 		},
 	}