1 files changed, 46 insertions, 4 deletions
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java
index fd282f9..798de28 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/output/lang/java/JavaWriterTests.java
@@ -7,6 +7,8 @@ import org.zwobble.hobgoblin.compiler.output.lang.java.ast.*;
 
 import java.io.IOException;
 import java.io.StringWriter;
+import java.io.Writer;
+import java.nio.file.NoSuchFileException;
 import java.util.List;
 import java.util.Optional;
 
@@ -15,6 +17,33 @@ import static org.zwobble.precisely.Matchers.equalTo;
 
 public class JavaWriterTests {
     @Test
+    public void compilationUnit() throws IOException {
+        var java = new JavaCompilationUnit(
+            List.of("com", "example", "shapes"),
+            new JavaClassDeclaration(
+                "Rectangle",
+                List.of(),
+                new JavaCustomArea(JavaTypeRef.topLevel(List.of(), "Rectangle"), "body"),
+                DocComment.EMPTY
+            )
+        );
+
+        var string = write(writer -> writer.writeCompilationUnit(java));
+
+        assertThat(string, equalTo("""
+            // Generated by hobgoblin.
+
+            package com.example.shapes;
+
+            public class Rectangle {
+                // Custom area start: Rectangle body
+                // Custom area end: Rectangle body
+            }
+            """
+        ));
+    }
+
+    @Test
     public void emptyClassDeclaration() throws IOException {
         var java = new JavaClassDeclaration(
             "Rectangle",
@@ -538,10 +567,23 @@ public class JavaWriterTests {
 
     private String write(Write write, Optional<String> originalContents) throws IOException {
         var stringWriter = new StringWriter();
-        var javaWriter = new JavaWriter(new CodeWriter(
-            stringWriter,
-            CodeWriter.CustomAreas.read(originalContents.orElse(""), "// Custom area start: ", "// Custom area end: ")
-        ));
+        var javaWriter = JavaWriter.file(
+            new CodeWriter.CodeWriterFile() {
+                @Override
+                public String readString() throws IOException {
+                    if (originalContents.isPresent()) {
+                        return originalContents.get();
+                    } else {
+                        throw new NoSuchFileException("");
+                    }
+                }
+
+                @Override
+                public Writer openWriter() throws IOException {
+                    return stringWriter;
+                }
+            }
+        );
         write.write(javaWriter);
         return stringWriter.toString();
     }