From a01da1a773c543da55b0878813d3a42818fe97da Mon Sep 17 00:00:00 2001
From: matthew16550 <matthew@matthew16550.com>
Date: Fri, 29 Oct 2021 15:36:37 +1100
Subject: [PATCH] Add more to PlantUmlTestUtils

---
 .../plantuml/test/PlantUmlTestUtils.java      | 22 ++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/test/net/sourceforge/plantuml/test/PlantUmlTestUtils.java b/test/net/sourceforge/plantuml/test/PlantUmlTestUtils.java
index f1d3e9cfc..e4bac23ce 100644
--- a/test/net/sourceforge/plantuml/test/PlantUmlTestUtils.java
+++ b/test/net/sourceforge/plantuml/test/PlantUmlTestUtils.java
@@ -1,11 +1,14 @@
 package net.sourceforge.plantuml.test;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
+import static org.assertj.core.api.Assertions.assertThat;
 
 import java.awt.image.BufferedImage;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.util.List;
 
 import net.sourceforge.plantuml.BlockUml;
@@ -28,6 +31,7 @@ public class PlantUmlTestUtils {
 		return new ExportDiagram(diagram);
 	}
 
+	@SuppressWarnings("unused")
 	public static class ExportDiagram {
 		private final Diagram diagram;
 		private boolean metadata;
@@ -36,6 +40,12 @@ public class PlantUmlTestUtils {
 			this.diagram = diagram;
 		}
 
+		public ExportDiagram assertDiagramType(Class<? extends Diagram> klass) {
+			assertNoError();
+			assertThat(diagram).isInstanceOf(klass);
+			return this;
+		}
+
 		public ExportDiagram assertNoError() {
 			if (diagram instanceof PSystemError) {
 				final PSystemError error = (PSystemError) this.diagram;
@@ -51,7 +61,11 @@ public class PlantUmlTestUtils {
 		}
 
 		public BufferedImage asImage() throws IOException {
-			return SImageIO.read(asByteArray(FileFormat.PNG));
+			return asImage(FileFormat.PNG);
+		}
+
+		public BufferedImage asImage(FileFormat fileFormat) throws IOException {
+			return SImageIO.read(asByteArray(fileFormat));
 		}
 
 		public String asString() throws IOException {
@@ -68,6 +82,12 @@ public class PlantUmlTestUtils {
 			return this;
 		}
 
+		public ExportDiagram toFile(Path path, FileFormat fileFormat) throws IOException {
+			try (OutputStream os = Files.newOutputStream(path)) {
+				return stream(os, fileFormat);
+			}
+		}
+
 		public ExportDiagram withMetadata(boolean metadata) {
 			this.metadata = metadata;
 			return this;