1+ #
2+ # Copyright (C) 2025 Google LLC
3+ #
4+ # Licensed under the Apache License, Version 2.0 (the "License"); you may not
5+ # use this file except in compliance with the License. You may obtain a copy of
6+ # the License at
7+ #
8+ # http://www.apache.org/licenses/LICENSE-2.0
9+ #
10+ # Unless required by applicable law or agreed to in writing, software
11+ # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
12+ # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
13+ # License for the specific language governing permissions and limitations under
14+ # the License.
15+ #
16+
17+ import logging
18+ import os
19+ import unittest
20+
21+ import apache_beam as beam
22+ import yaml
23+ from apache_beam .yaml import yaml_transform
24+ from jinja2 import Environment , FileSystemLoader , meta
25+
26+ # Configure logging at the module level to ensure it's set up early.
27+ # This will ensure logs are printed to console.
28+ # Using basicConfig with a format to make logs more informative.
29+ logging .basicConfig (level = logging .DEBUG ,
30+ format = '%(asctime)s - %(name)s - %(levelname)s - %(message)s' )
31+ logging .getLogger ('apache_beam' ).setLevel (logging .DEBUG )
32+
33+
34+ def create_test_method (template_name , yaml_dir ):
35+ """Creates a test method that validates a single YAML template.
36+
37+ This factory function generates a test method that will be dynamically added
38+ to the YamlSyntaxTest class. Each generated test validates a single YAML
39+ template file by rendering it with placeholder values for any Jinja
40+ variables and then validating the resulting YAML against the Beam 'generic'
41+ schema.
42+
43+ Args:
44+ template_name (str): The filename of the YAML template to be tested.
45+ yaml_dir (str): The directory where the YAML templates are located.
46+
47+ Returns:
48+ function: A test method that can be attached to a unittest.TestCase class.
49+ """
50+
51+ def test_method (self ):
52+ self ._logger .info (f"Validating { template_name } )
53+ env = Environment (loader = FileSystemLoader (yaml_dir ), autoescape = False )
54+ template_source = env .loader .get_source (env , template_name )[0 ]
55+
56+ # Find all undeclared variables in the template
57+ parsed_content = env .parse (template_source )
58+ undeclared_vars = meta .find_undeclared_variables (parsed_content )
59+
60+ # Use placeholder values for Jinja variables for validation purposes
61+ context = {var : 'placeholder' for var in undeclared_vars }
62+ template = env .get_template (template_name )
63+ rendered_yaml = template .render (context )
64+ self ._logger .debug (f"Rendered YAML for { template_name } \n { rendered_yaml } )
65+
66+ self ._logger .debug (f"Loading YAML into Beam pipeline_spec: { template_name } )
67+ pipeline_spec = yaml .load (rendered_yaml , Loader = yaml_transform .SafeLineLoader )
68+
69+ # Validate the pipeline spec against the generic schema without trying to
70+ # expand the transforms, which avoids the need for expansion services.
71+ yaml_transform .validate_against_schema (pipeline_spec , 'generic' )
72+ self ._logger .info (f"Successfully validated YAML syntax for: { template_name } )
73+
74+ return test_method
75+
76+
77+ class YamlSyntaxTest (unittest .TestCase ):
78+ """A test suite for validating the syntax of Beam YAML templates.
79+
80+ This class is dynamically populated with test methods, one for each
81+ .yaml file found in the `src/main/yaml` directory. This is accomplished
82+ by the `_create_tests` function, which runs at module-load time.
83+ """
84+ _logger = logging .getLogger (__name__ )
85+
86+
87+ def _create_tests ():
88+ """Discovers all YAML templates and dynamically creates a test for each.
89+
90+ This function scans the `src/main/yaml` directory for `.yaml` files and,
91+ for each file, generates a unique test method on the `YamlSyntaxTest` class.
92+ This allows `unittest` or `pytest` to discover and run each validation as a
93+ separate test case, making it easy to identify which template is invalid.
94+ """
95+ yaml_dir = os .path .join (os .path .dirname (__file__ ), '../../main/yaml' )
96+ if not os .path .isdir (yaml_dir ):
97+ return
98+
99+ env = Environment (loader = FileSystemLoader (yaml_dir ))
100+ for template_name in env .list_templates (filter_func = lambda x : x .endswith ('.yaml' )):
101+ test_name = f"test_{ template_name .replace ('.yaml' , '' ).replace ('-' , '_' )}
102+ test_method = create_test_method (template_name , yaml_dir )
103+ setattr (YamlSyntaxTest , test_name , test_method )
104+
105+ _create_tests ()
106+
107+ if __name__ == '__main__' :
108+ unittest .main ()
0 commit comments