Centralize the template type resolution logic

Bug: 330885075
Change-Id: I5b3080f730b28a80791554a4c4e48156213ae8df
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/5376633
Reviewed-by: Andres Ricardo Perez <andresrperez@chromium.org>
Commit-Queue: Jean-Philippe Gravel <jpgravel@chromium.org>
Cr-Commit-Position: refs/heads/main@{#1282136}
diff --git a/html/canvas/tools/gentestutilsunion.py b/html/canvas/tools/gentestutilsunion.py
index 3aad13b..9df6839 100644
--- a/html/canvas/tools/gentestutilsunion.py
+++ b/html/canvas/tools/gentestutilsunion.py
@@ -253,6 +253,12 @@
     WORKER = 'Worker'
 
 
+class _TemplateType(str, enum.Enum):
+    REFERENCE = 'reference'
+    HTML_REFERENCE = 'html_reference'
+    TESTHARNESS = 'testharness'
+
+
 @dataclasses.dataclass
 class _OutputPaths:
     element: str
@@ -281,11 +287,6 @@
             '"promise" now for creating promise test type in the template '
             'file.')
 
-    if 'reference' in test and 'html_reference' in test:
-        raise InvalidTestDefinitionError(
-            f'Test {test["name"]} is invalid, "reference" and "html_reference" '
-            'can\'t both be specified at the same time.')
-
 
 def _render_template(jinja_env: jinja2.Environment, template: jinja2.Template,
                      params: _TestParams) -> str:
@@ -359,11 +360,25 @@
     return frozenset(_CanvasType(t) for t in canvas_types)
 
 
+def _get_template_type(params: _TestParams) -> _TemplateType:
+    if 'reference' in params and 'html_reference' in params:
+        raise InvalidTestDefinitionError(
+            f'Test {params["name"]} is invalid, "reference" and '
+            '"html_reference" can\'t both be specified at the same time.')
+
+    if 'reference' in params:
+        return _TemplateType.REFERENCE
+    if 'html_reference' in params:
+        return _TemplateType.HTML_REFERENCE
+    return _TemplateType.TESTHARNESS
+
+
 def _finalize_params(jinja_env: jinja2.Environment,
                      params: _MutableTestParams) -> None:
     params['name'] = _get_variant_name(jinja_env, params)
     params['file_name'] = _get_file_name(params)
     params['canvas_types'] = _get_canvas_types(params)
+    params['template_type'] = _get_template_type(params)
 
 
 def _write_reference_test(jinja_env: jinja2.Environment, params: _TestParams,
@@ -480,8 +495,8 @@
     })
 
     output_files = output_dirs.sub_path(params['file_name'])
-
-    if 'reference' in test or 'html_reference' in test:
+    if params['template_type'] in (_TemplateType.REFERENCE,
+                                   _TemplateType.HTML_REFERENCE):
         _write_reference_test(jinja_env, params, output_files)
     else:
         _write_testharness_test(jinja_env, params, output_files)