diff --git a/diagram/src/test/java/io/serverlessworkflow/diagram/test/CustomTemplateWorkflowDiagramTest.java b/diagram/src/test/java/io/serverlessworkflow/diagram/test/CustomTemplateWorkflowDiagramTest.java
index 530d4b8c..5c9b7038 100644
--- a/diagram/src/test/java/io/serverlessworkflow/diagram/test/CustomTemplateWorkflowDiagramTest.java
+++ b/diagram/src/test/java/io/serverlessworkflow/diagram/test/CustomTemplateWorkflowDiagramTest.java
@@ -39,12 +39,15 @@ public void testSpecExamplesParsing(String workflowLocation) throws Exception {
     assertNotNull(workflow.getStates());
 
     WorkflowDiagram workflowDiagram =
-        new WorkflowDiagramImpl().setWorkflow(workflow).setTemplate("custom-template");
+        new WorkflowDiagramImpl()
+            .showLegend(true)
+            .setWorkflow(workflow)
+            .setTemplate("custom-template");
 
     String diagramSVG = workflowDiagram.getSvgDiagram();
 
     Assertions.assertNotNull(diagramSVG);
-    // custom template uses #0000FF as start node color
-    Assertions.assertTrue(diagramSVG.contains("#0000FF"));
+    // custom template uses customcolor in the legend
+    Assertions.assertTrue(diagramSVG.contains("customcolor"));
   }
 }
diff --git a/diagram/src/test/resources/templates/plantuml/custom-template.txt b/diagram/src/test/resources/templates/plantuml/custom-template.txt
index c1794f55..e162afc5 100644
--- a/diagram/src/test/resources/templates/plantuml/custom-template.txt
+++ b/diagram/src/test/resources/templates/plantuml/custom-template.txt
@@ -3,7 +3,7 @@ skinparam backgroundColor White
 skinparam legendBackgroundColor White
 skinparam legendBorderColor White
 skinparam state {
-  StartColor #0000FF
+  StartColor Green
   EndColor Orange
   BackgroundColor GhostWhite
   BackgroundColor<< workflow >> White
@@ -39,7 +39,7 @@ state "[(${diagram.title})]" as workflow << workflow >> {
 legend center
 State Types and Border Colors:
 | Event | Operation | Switch | Sleep | Parallel | Inject | ForEach | CallBack |
-|<#7fe5f0>|<#bada55>|<#92a0f2>|<#b83b5e>|<#6a2c70>|<#1e5f74>|<#931a25>|<#ffcb8e>|
+|<#7fe5f0>|<#bada55>|<#92a0f2>|<#b83b5e>|<#6a2c70>|<#1e5f74>|<#931a25>|<customcolor>|
 endlegend
 [/]