1 files changed, 10 insertions, 1 deletions
diff --git a/src/partest/scala/tools/partest/ScaladocModelTest.scala b/src/partest/scala/tools/partest/ScaladocModelTest.scala
index fb93e98726..fc540f83fe 100644
--- a/src/partest/scala/tools/partest/ScaladocModelTest.scala
+++ b/src/partest/scala/tools/partest/ScaladocModelTest.scala
@@ -11,8 +11,9 @@ import scala.tools.nsc._
 import scala.tools.nsc.util.CommandLineParser
 import scala.tools.nsc.doc.{Settings, DocFactory, Universe}
 import scala.tools.nsc.doc.model._
-import scala.tools.nsc.reporters.ConsoleReporter
+import scala.tools.nsc.doc.model.diagram._
 import scala.tools.nsc.doc.model.comment._
+import scala.tools.nsc.reporters.ConsoleReporter
 
 /** A class for testing scaladoc model generation
  *   - you need to specify the code in the `code` method
@@ -181,5 +182,13 @@ abstract class ScaladocModelTest extends DirectTest {
       }
       countLinks(c.body)
     }
+
+    def testDiagram(doc: DocTemplateEntity, diag: Option[Diagram], nodes: Int, edges: Int) = {
+      assert(diag.isDefined, doc.qualifiedName + " diagram missing")
+      assert(diag.get.nodes.length == nodes,
+             doc.qualifiedName + "'s diagram: node count " + diag.get.nodes.length + " == " + nodes)
+      assert(diag.get.edges.map(_._2.length).sum == edges,
+             doc.qualifiedName + "'s diagram: edge count " + diag.get.edges.length + " == " + edges)
+    }
   }
 }