1 files changed, 33 insertions, 26 deletions
diff --git a/testsuite/Testsrc/Testlib/TestServer/TestPlugins/TestCfg/TestCfgCheetahGenerator.py b/testsuite/Testsrc/Testlib/TestServer/TestPlugins/TestCfg/TestCfgCheetahGenerator.py
index fc5d5e53d..93331304a 100644
--- a/testsuite/Testsrc/Testlib/TestServer/TestPlugins/TestCfg/TestCfgCheetahGenerator.py
+++ b/testsuite/Testsrc/Testlib/TestServer/TestPlugins/TestCfg/TestCfgCheetahGenerator.py
@@ -17,32 +17,39 @@ from common import *
 from TestServer.TestPlugins.TestCfg.Test_init import TestCfgGenerator
 
 
-if HAS_CHEETAH or can_skip:
-    class TestCfgCheetahGenerator(TestCfgGenerator):
-        test_obj = CfgCheetahGenerator
+class TestCfgCheetahGenerator(TestCfgGenerator):
+    test_obj = CfgCheetahGenerator
 
-        @skipUnless(HAS_CHEETAH, "Cheetah libraries not found, skipping")
-        def setUp(self):
-            pass
+    @skipUnless(HAS_CHEETAH, "Cheetah libraries not found, skipping")
+    def setUp(self):
+        TestCfgGenerator.setUp(self)
+        set_setup_default("repository", datastore)
 
-        @patch("Bcfg2.Server.Plugins.Cfg.CfgCheetahGenerator.Template")
-        def test_get_data(self, mock_Template):
-            ccg = self.get_obj(encoding='UTF-8')
-            ccg.data = "data"
-            entry = lxml.etree.Element("Path", name="/test.txt")
-            metadata = Mock()
-            Bcfg2.Server.Plugins.Cfg.CfgCheetahGenerator.SETUP = MagicMock()
+    @patch("Bcfg2.Server.Plugins.Cfg.CfgCheetahGenerator.Template")
+    @patch("Bcfg2.Server.Plugins.Cfg.CfgCheetahGenerator.get_template_data")
+    def test_get_data(self, mock_get_template_data, mock_Template):
+        ccg = self.get_obj()
+        ccg.data = "data"
+        entry = lxml.etree.Element("Path", name="/test.txt")
+        metadata = Mock()
 
-            self.assertEqual(ccg.get_data(entry, metadata),
-                             mock_Template.return_value.respond.return_value)
-            Bcfg2.Server.Plugins.Cfg.CfgCheetahGenerator.SETUP.__getitem__.assert_called_with("repo")
-            mock_Template.assert_called_with("data".decode(ccg.encoding),
-                                             compilerSettings=ccg.settings)
-            tmpl = mock_Template.return_value
-            tmpl.respond.assert_called_with()
-            self.assertEqual(tmpl.metadata, metadata)
-            self.assertEqual(tmpl.name, entry.get("name"))
-            self.assertEqual(tmpl.path, entry.get("name"))
-            self.assertEqual(tmpl.source_path, ccg.name)
-            self.assertEqual(tmpl.repo,
-                             Bcfg2.Server.Plugins.Cfg.CfgCheetahGenerator.SETUP.__getitem__.return_value)
+        template_vars = dict(name=entry.get("name"),
+                             metadata=metadata,
+                             path=ccg.name,
+                             source_path=ccg.name,
+                             repo=datastore)
+        mock_get_template_data.return_value = template_vars
+
+        self.assertEqual(ccg.get_data(entry, metadata),
+                         mock_Template.return_value.respond.return_value)
+        mock_Template.assert_called_with(
+            "data".decode(Bcfg2.Options.setup.encoding),
+            compilerSettings=ccg.settings)
+        tmpl = mock_Template.return_value
+        tmpl.respond.assert_called_with()
+        for key, val in template_vars.items():
+            self.assertEqual(getattr(tmpl, key), val)
+        self.assertItemsEqual(mock_get_template_data.call_args[0],
+                              [entry, metadata, ccg.name])
+        self.assertIsInstance(mock_get_template_data.call_args[1]['default'],
+                              DefaultCheetahDataProvider)