diff --git a/asciidoctorj-diagram/gradle.properties b/asciidoctorj-diagram/gradle.properties
index 90ed7d6..8a0a414 100644
--- a/asciidoctorj-diagram/gradle.properties
+++ b/asciidoctorj-diagram/gradle.properties
@@ -1,4 +1,4 @@
 properName=AsciidoctorJ Diagram
 description=AsciidoctorJ Diagram bundles the Asciidoctor Diagram RubyGem (asciidoctor-diagram) so it can be loaded into the JVM using JRuby.
-version=2.0.4
+version=2.0.5
 gem_name=asciidoctor-diagram
diff --git a/asciidoctorj-diagram/src/test/java/org.asciidoctor/WhenDocumentContainsDitaaDiagram.java b/asciidoctorj-diagram/src/test/java/org.asciidoctor/WhenDocumentContainsDitaaDiagram.java
index 1b8d5a3..c9bb2d6 100644
--- a/asciidoctorj-diagram/src/test/java/org.asciidoctor/WhenDocumentContainsDitaaDiagram.java
+++ b/asciidoctorj-diagram/src/test/java/org.asciidoctor/WhenDocumentContainsDitaaDiagram.java
@@ -15,16 +15,16 @@ public class WhenDocumentContainsDitaaDiagram {
     @Test
     public void png_should_be_rendered_for_diagram() {
 
-        File inputFile = new File("build/resources/test/sample.adoc");
-        File outputDir = new File("build");
-        File outputFile1 = new File(outputDir, "asciidoctor-diagram-process.png");
-        File outputFile2 = new File(outputDir, ".asciidoctor/diagram/asciidoctor-diagram-process.png.cache");
+        File buildDir = new File("build/resources/test");
+
+        File inputFile = new File(buildDir, "sample.adoc");
+        File outputFile1 = new File(inputFile.getParentFile(), "asciidoctor-diagram-process.png");
+        File outputFile2 = new File(inputFile.getParentFile(), ".asciidoctor/diagram/asciidoctor-diagram-process.png.cache");
         asciidoctor.requireLibrary("asciidoctor-diagram");
         asciidoctor.convertFile(inputFile,
             options()
                 .backend("html5")
-                .toDir(outputDir)
-                .safe(SafeMode.UNSAFE)
+                .toFile(new File(buildDir, "sample.html"))
                 .get());
         assertThat(outputFile1.exists(), is(true));
         assertThat(outputFile2.exists(), is(true));
diff --git a/asciidoctorj-diagram/src/test/java/org.asciidoctor/WhenDocumentWantsDataUris.java b/asciidoctorj-diagram/src/test/java/org.asciidoctor/WhenDocumentWantsDataUris.java
index 4f9df55..07f201d 100644
--- a/asciidoctorj-diagram/src/test/java/org.asciidoctor/WhenDocumentWantsDataUris.java
+++ b/asciidoctorj-diagram/src/test/java/org.asciidoctor/WhenDocumentWantsDataUris.java
@@ -6,6 +6,9 @@
 import org.junit.Test;
 
 import java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
 
 import static org.asciidoctor.OptionsBuilder.options;
 import static org.hamcrest.CoreMatchers.is;
@@ -18,16 +21,17 @@ public class WhenDocumentWantsDataUris {
     private Asciidoctor asciidoctor = Asciidoctor.Factory.create();
 
     @Test
-    public void png_should_be_rendered_for_diagram() {
+    public void png_should_be_rendered_for_diagram() throws IOException {
+        File buildDir = new File("build/resources/test");
 
-        File inputFile = new File("build/resources/test/data-uri.adoc");
+        File inputFile = new File(buildDir, "data-uri.adoc");
+        File outputFile = new File(buildDir, "data-uri.html");
         File outputFile1 = new File(inputFile.getParentFile(), "data-uri-test.png");
         File outputFile2 = new File(inputFile.getParentFile(), ".asciidoctor/diagram/data-uri-test.png.cache");
         asciidoctor.requireLibrary("asciidoctor-diagram");
-        final String html = asciidoctor.convertFile(inputFile,
+        asciidoctor.convertFile(inputFile,
             options().backend("html5")
-                .toFile(false)
-                .toDir(new File("build"))
+                .toFile(outputFile)
                 .safe(SafeMode.SERVER)
                 .get());
         assertThat(outputFile1.exists(), is(true));
@@ -35,6 +39,7 @@ public void png_should_be_rendered_for_diagram() {
         outputFile1.delete();
         outputFile2.delete();
 
+        String html = new String(Files.readAllBytes(outputFile.toPath()), StandardCharsets.UTF_8);
         final Document doc = Jsoup.parse(html);
         System.out.println(doc);
         Elements images = doc.getElementsByTag("img");
diff --git a/gradle.properties b/gradle.properties
index 6d986d8..0818d8f 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -1 +1 @@
-version=2.0.4
+version=2.0.5
\ No newline at end of file