diff --git a/jte-models/src/test/java/gg/jte/models/generator/TestModelExtension.java b/jte-models/src/test/java/gg/jte/models/generator/TestModelExtension.java
index 5840dc64..cc00b400 100644
--- a/jte-models/src/test/java/gg/jte/models/generator/TestModelExtension.java
+++ b/jte-models/src/test/java/gg/jte/models/generator/TestModelExtension.java
@@ -249,6 +249,7 @@ fun hello(content: gg.jte.Content): JteModel
 
                     import gg.jte.models.runtime.*
                     import gg.jte.ContentType
+                    import gg.jte.TemplateOutput
                     import gg.jte.html.HtmlTemplateOutput
 
                     class StaticTemplates : Templates {
@@ -331,6 +332,7 @@ fun hello(): JteModel
 
                     import gg.jte.models.runtime.*
                     import gg.jte.ContentType
+                    import gg.jte.TemplateOutput
                     import gg.jte.html.HtmlTemplateOutput
 
                     class StaticTemplates : Templates {