Improve jinja2 search, include custom config path into search

Signed-off-by: Denys Fedoryshchenko <[email protected]>

Author: nuclearcat

kernelci/cli/job.py

@@ -55,7 +55,7 @@ def new(name, input_node_id, platform,  # pylint: disable=too-many-arguments
         if runtime not in configs['runtimes']:
             raise click.ClickException(f"Invalid runtime {runtime}")
         runtime = kernelci.runtime.get_runtime(
-            configs['runtimes'][runtime], token=secrets.api.runtime_token)
+            configs['runtimes'][runtime], token=secrets.api.runtime_token, custom_template_dir=config[0] if config else None)
     job_node = helper.create_job_node(job_config, input_node,
                                       platform=platform_config, runtime=runtime)
     if job_node:
@@ -116,7 +116,7 @@ def generate(node_id,  # pylint: disable=too-many-arguments, too-many-locals
     if runtime_config is None:
         raise click.ClickException(f"Runtime {runtime} not found in the config")
     runtime = kernelci.runtime.get_runtime(
-        runtime_config, token=secrets.api.runtime_token)
+        runtime_config, token=secrets.api.runtime_token, custom_template_dir=config[0] if config else None)
     params = runtime.get_params(job, api.config)
     if not params:
         raise click.ClickException("Invalid job parameters, aborting...")
@@ -149,7 +149,7 @@ def submit(runtime, job_path, wait,  # pylint: disable=too-many-arguments
     configs = kernelci.config.load(config)
     runtime_config = configs['runtimes'][runtime]
     runtime = kernelci.runtime.get_runtime(
-        runtime_config, token=secrets.api.runtime_token
+        runtime_config, token=secrets.api.runtime_token, custom_template_dir=config[0] if config else None
     )
     job = runtime.submit(job_path)
     click.echo(runtime.get_job_id(job))

kernelci/runtime/__init__.py

@@ -12,7 +12,7 @@
 import requests
 import yaml
 
-from jinja2 import Environment, FileSystemLoader
+from jinja2 import Environment, FileSystemLoader, ChoiceLoader
 
 from kernelci.config.base import get_system_arch
 
@@ -75,13 +75,22 @@ class Runtime(abc.ABC):
     TEMPLATES = ['config/runtime', '/etc/kernelci/runtime']
 
     # pylint: disable=unused-argument
-    def __init__(self, config, user=None, token=None):
+    def __init__(self, config, user=None, token=None, custom_template_dir=None):
         """A Runtime object can be used to run jobs in a runtime environment
 
         *config* is a kernelci.config.runtime.Runtime object
+        *custom_template_dir* is an optional custom directory for Jinja2 templates
         """
         self._config = config
-        self._templates = self.TEMPLATES
+        self._templates = self.TEMPLATES.copy()
+        if custom_template_dir:
+            # Add the main custom dir
+            self._templates.append(custom_template_dir)
+            # Add relevant subdirectories
+            for subdir in ["runtime", "runtime/base", "runtime/boot", "runtime/tests"]:
+                sub_path = os.path.join(custom_template_dir, subdir) if not custom_template_dir.endswith(subdir) else custom_template_dir
+                if os.path.isdir(sub_path):
+                    self._templates.append(sub_path)
         self._user = user
         self._token = token
 
@@ -96,8 +105,9 @@ def templates(self):
         return self._templates
 
     def _get_template(self, job_config):
+        loaders = [FileSystemLoader(path) for path in self.templates]
         jinja2_env = Environment(
-            loader=FileSystemLoader(self.templates),
+            loader=ChoiceLoader(loaders),
             extensions=["jinja2.ext.do"]
         )
         jinja2_env.globals.update(self._get_jinja2_functions())
@@ -211,19 +221,20 @@ def wait(self, job_object):
         """Wait for a job to complete and get the exit status code"""
 
 
-def get_runtime(config, user=None, token=None):
+def get_runtime(config, user=None, token=None, custom_template_dir=None):
     """Get the Runtime object for a given runtime config.
 
     *config* is a kernelci.config.runtime.Runtime object
     *user* is the name of the user to connect to the runtime
     *token* is the associated token to connect to the runtime
+    *custom_template_dir* is an optional custom directory for Jinja2 templates
     """
     module_name = '.'.join(['kernelci', 'runtime', config.lab_type])
     runtime_module = importlib.import_module(module_name)
-    return runtime_module.get_runtime(config, user=user, token=token)
+    return runtime_module.get_runtime(config, user=user, token=token, custom_template_dir=custom_template_dir)
 
 
-def get_all_runtimes(runtime_configs, opts):
+def get_all_runtimes(runtime_configs, opts, custom_template_dir=None):
     """Get all the Runtime objects based on the runtime configs and options
 
     This will iterate over all the runtimes configs and yield a (name, runtime)
@@ -232,14 +243,16 @@ def get_all_runtimes(runtime_configs, opts):
 
     *runtime_configs* is the 'runtimes' config loaded from YAML
     *opts* is an Options object loaded from the CLI args and settings file
+    *custom_template_dir* is an optional custom directory for Jinja2 templates
     """
     for config_name, config in runtime_configs.items():
         section = ('runtime', config_name)
         user, token = (
             opts.get_from_section(section, opt)
             for opt in ('user', 'runtime_token')
         )
-        yield config_name, get_runtime(config, user, token)
+        runtime = get_runtime(config, user=user, token=token, custom_template_dir=custom_template_dir)
+        yield config_name, runtime
 
 
 def evaluate_test_suite_result(child_nodes):