diff --git a/java/java_test.go b/java/java_test.go
index 8bae019..5d6a6e0 100644
--- a/java/java_test.go
+++ b/java/java_test.go
@@ -81,6 +81,7 @@
 	ctx.RegisterModuleType("genrule", android.ModuleFactoryAdaptor(genrule.GenRuleFactory))
 	ctx.RegisterModuleType("droiddoc", android.ModuleFactoryAdaptor(DroiddocFactory))
 	ctx.RegisterModuleType("droiddoc_host", android.ModuleFactoryAdaptor(DroiddocHostFactory))
+	ctx.RegisterModuleType("droiddoc_template", android.ModuleFactoryAdaptor(DroiddocTemplateFactory))
 	ctx.PreArchMutators(android.RegisterPrebuiltsPreArchMutators)
 	ctx.PreArchMutators(android.RegisterPrebuiltsPostDepsMutators)
 	ctx.PreArchMutators(android.RegisterDefaultsPreArchMutators)
@@ -865,6 +866,10 @@
 
 func TestDroiddoc(t *testing.T) {
 	ctx := testJava(t, `
+		droiddoc_template {
+		    name: "droiddoc-templates-sdk",
+		    path: ".",
+		}
 		droiddoc {
 		    name: "bar-doc",
 		    srcs: [
@@ -873,7 +878,7 @@
 		    exclude_srcs: [
 		        "bar-doc/b.java"
 		    ],
-		    custom_template_dir: "external/doclava/templates-sdk",
+		    custom_template: "droiddoc-templates-sdk",
 		    hdf: [
 		        "android.whichdoc offline",
 		    ],