load_template searchpaths, take 2

napalm_base/helpers.py

@@ -37,29 +37,28 @@ class _MACFormat(mac_unix):
 def load_template(cls, template_name, template_source=None, template_path=None,
                   openconfig=False, **template_vars):
     try:
+        search_path = []
         if isinstance(template_source, py23_compat.string_types):
             template = jinja2.Template(template_source)
         else:
-            current_dir = os.path.dirname(os.path.abspath(sys.modules[cls.__module__].__file__))
-            if (isinstance(template_path, py23_compat.string_types) and
-                    os.path.isdir(template_path) and os.path.isabs(template_path)):
-                current_dir = os.path.join(template_path, cls.__module__.split('.')[-1])
-                # append driver name at the end of the custom path
+            if template_path is not None:
+                if (isinstance(template_path, py23_compat.string_types) and
+                        os.path.isdir(template_path) and os.path.isabs(template_path)):
+                    # append driver name at the end of the custom path
+                    search_path.append(os.path.join(template_path, cls.__module__.split('.')[-1]))
+                else:
+                    raise IOError("Template path does not exist: {}".format(template_path))
+            else:
+                # Search modules for template paths
+                search_path = [os.path.dirname(os.path.abspath(sys.modules[c.__module__].__file__))
+                               for c in cls.__class__.mro() if c is not object]
 
             if openconfig:
-                template_dir_path = '{current_dir}/oc_templates'.format(current_dir=current_dir)
+                search_path = ['{}/oc_templates'.format(s) for s in search_path]
             else:
-                template_dir_path = '{current_dir}/templates'.format(current_dir=current_dir)
-
-            if not os.path.isdir(template_dir_path):
-                raise napalm_base.exceptions.DriverTemplateNotImplemented(
-                        '''Config template dir does not exist: {path}.
-                        Please create it and add driver-specific templates.'''.format(
-                            path=template_dir_path
-                        )
-                    )
+                search_path = ['{}/templates'.format(s) for s in search_path]
 
-            loader = jinja2.FileSystemLoader(template_dir_path)
+            loader = jinja2.FileSystemLoader(search_path)
             environment = jinja2.Environment(loader=loader)
 
             for filter_name, filter_function in CustomJinjaFilters.filters().items():
@@ -71,9 +70,9 @@ def load_template(cls, template_name, template_source=None, template_path=None,
         configuration = template.render(**template_vars)
     except jinja2.exceptions.TemplateNotFound:
         raise napalm_base.exceptions.TemplateNotImplemented(
-            "Config template {template_name}.j2 is not defined under {path}".format(
+            "Config template {template_name}.j2 not found in search path: {sp}".format(
                 template_name=template_name,
-                path=template_dir_path
+                sp=search_path
             )
         )
     except (jinja2.exceptions.UndefinedError, jinja2.exceptions.TemplateSyntaxError) as jinjaerr:

test/unit/test_helpers.py

@@ -59,12 +59,14 @@ def test_load_template(self):
             * check if can load empty template
             * check if raises TemplateRenderException when template is not correctly formatted
             * check if can load correct template
-            * check if can load correct template even if wrong custom path specified
+            * check if raises IOError if invalid path is specified
             * check if raises TemplateNotImplemented when trying to use inexisting template in
               custom path
             * check if can load correct template from custom path
             * check if template passed as string can be loaded
+            * check that the search path setup by MRO is correct when loading an incorrecet template
         """
+
         self.assertTrue(HAS_JINJA)  # firstly check if jinja2 is installed
         _NTP_PEERS_LIST = [
             '172.17.17.1',
@@ -94,14 +96,16 @@ def test_load_template(self):
                                                           '__a_very_nice_template__',
                                                           **_TEMPLATE_VARS))
 
-        self.assertTrue(napalm_base.helpers.load_template(self.network_driver,
-                                                          '__a_very_nice_template__',
-                                                          template_path='/this/path/does/not/exist',
-                                                          **_TEMPLATE_VARS))
+        self.assertRaises(IOError,
+                          napalm_base.helpers.load_template,
+                          self.network_driver,
+                          '__a_very_nice_template__',
+                          template_path='/this/path/does/not/exist',
+                          **_TEMPLATE_VARS)
 
         install_dir = os.path.dirname(
             os.path.abspath(sys.modules[self.network_driver.__module__].__file__))
-        custom_path = os.path.join(install_dir, 'test/custom/path/base')
+        custom_path = os.path.join(install_dir, '../custom/path/base')
 
         self.assertRaises(napalm_base.exceptions.TemplateNotImplemented,
                           napalm_base.helpers.load_template,
@@ -121,6 +125,13 @@ def test_load_template(self):
                                                           '_this_still_needs_a_name',
                                                           template_source=template_source,
                                                           **_TEMPLATE_VARS))
+        self.assertRaisesRegexp(napalm_base.exceptions.TemplateNotImplemented,
+                                "path.*napalm-base/test/unit/templates'"
+                                + ",.*napalm-base/napalm_base/templates']",
+                                napalm_base.helpers.load_template,
+                                self.network_driver,
+                                '__this_template_does_not_exist__',
+                                **_TEMPLATE_VARS)
 
     def test_textfsm_extractor(self):
         """