Allow environment variables in crontab file

.gitignore

@@ -2,3 +2,4 @@ build
 *.pyc
 dist
 django_kronos.egg-info
+.idea/

CONTRIBUTORS.txt

@@ -1,3 +1,4 @@
 Rodrigo N. Carreras (https://github.com/carrerasrodrigo)
 Joshua Blum (https://github.com/joshblum)
-Paulo Cheadi Haddad Filho (https://github.com/paulochf)
+Paulo Cheadi Haddad Filho (https://github.com/paulochf)
+Russ Ferriday (https://github.com/topiaruss)

kronos/__init__.py

@@ -9,7 +9,7 @@
     from django.utils.importlib import import_module
 
 from kronos.settings import PROJECT_MODULE, KRONOS_PYTHON, KRONOS_MANAGE, \
-    KRONOS_PYTHONPATH, KRONOS_POSTFIX, KRONOS_PREFIX
+    KRONOS_PYTHONPATH, KRONOS_POSTFIX, KRONOS_PREFIX, KRONOS_ENV
 from django.conf import settings
 from kronos.utils import read_crontab, write_crontab, delete_crontab
 from kronos.version import __version__
@@ -119,6 +119,10 @@ def install():
     current_crontab = six.u(read_crontab())
 
     new_crontab = ''
+
+    for env_line in KRONOS_ENV.splitlines():
+        new_crontab += '# KRONOS_ENV_BREAD_CRUMB for next\n%s\n' % env_line
+
     for task in tasks:
         new_crontab += '%s\n' % task['fn'].cron_expression
 
@@ -137,15 +141,27 @@ def printtasks():
 
 
 def find_existing_jobs(current_crontab):
-    new_crontab = ''
-    for line in six.u(current_crontab).split('\n')[:-1]:
-        exp = '%(python)s %(manage)s runtask' % {
-            'python': KRONOS_PYTHON,
-            'manage': KRONOS_MANAGE,
-            }
+    remaining = []
+    exp = '%(python)s %(manage)s runtask' % {
+        'python': KRONOS_PYTHON,
+        'manage': KRONOS_MANAGE,
+        }
+    for line in six.u(current_crontab).splitlines():
         if not ('$KRONOS_BREAD_CRUMB' in line and exp in line):
-            new_crontab += '%s\n' % line
-    return new_crontab
+             remaining.append(line)
+
+    without_env = []
+    skip = False
+    for line in remaining:
+        if 'KRONOS_ENV_BREAD_CRUMB' in line:
+            skip = True
+            continue
+        if skip:
+            skip = False
+            continue
+        without_env += [line]
+
+    return "\n".join(without_env)
 
 
 def uninstall():

kronos/settings.py

@@ -9,3 +9,4 @@
 PROJECT_MODULE = sys.modules['.'.join(settings.SETTINGS_MODULE.split('.')[:-1])]
 KRONOS_POSTFIX = getattr(settings, 'KRONOS_POSTFIX', '')
 KRONOS_PREFIX = getattr(settings, 'KRONOS_PREFIX', '')
+KRONOS_ENV = '\n'.join(getattr(settings, 'KRONOS_ENV', '').split('\\n'))

kronos/tests/tests.py

@@ -23,7 +23,7 @@ def setUp(self):
         load()
 
     @patch('subprocess.Popen')
-    def test_unintalltasks(self, mock):
+    def test_uninstalltasks(self, mock):
         """Test uninstalling tasks with the ``uninstalltasks`` command."""
         mock.return_value = Mock(
             stdout=StringIO('crontab: installing new crontab'),
@@ -35,17 +35,28 @@ def test_unintalltasks(self, mock):
 
     def test_find_existing_jobs(self):
         """Test uninstalling tasks with the ``uninstalltasks`` command."""
+        env = """EXISTING1=other1
+        # KRONOS_ENV_BREAD_CRUMB for next line
+        INSERTED_BY_KRONOS=important
+        EXISTING2=other2
+        """
         keep = '%(python)s %(manage)s runtask' % {
             'python': KRONOS_PYTHON,
             'manage': KRONOS_MANAGE,
         }
         keep2 = " keep_me $KRONOS_BREAD_CRUMB"
         remove = keep + keep2
-        new_cron = find_existing_jobs("\n".join([keep, keep2, remove, ""]))
+        new_cron = find_existing_jobs(env + "\n".join([
+            keep, keep2, remove, ""]))
         self.assertIn(keep, new_cron)
         self.assertIn(keep2, new_cron)
         self.assertNotIn(remove, new_cron)
 
+        self.assertIn('EXISTING1', new_cron)
+        self.assertIn('EXISTING2', new_cron)
+        self.assertNotIn('KRONOS_ENV_BREAD_CRUMB', new_cron)
+        self.assertNotIn('INSERTED_BY_KRONOS', new_cron)
+
 
     @patch('subprocess.Popen')
     def test_read_crontab(self, mock):
@@ -134,16 +145,18 @@ def test_installtasks(self, mock):
 
         self.assertTrue(mock.called)
 
+    @patch('kronos.KRONOS_ENV', """FOO=bar\nBLACK=white""")
     @patch('subprocess.Popen')
     def test_installed_tasks(self, mock):
         """Test installing tasks with the ``installtasks`` command."""
         mock.return_value = Mock(
             stdout=StringIO('crontab: installing new crontab'),
             stderr=StringIO('')
         )
-
         call_command('installtasks')
         calls = str(mock.mock_calls[-1])
+        self.assertIn('FOO=bar', calls)
+        self.assertIn('runtask praise', calls)
         self.assertIn('runtask praise', calls)
         self.assertIn('runtask complain', calls)
         self.assertIn('manage.py task', calls)