Merge pull request #53 from moshez/environment-inheritance

Allow partial inheritance of the environment
ncolony · Sep 19, 2017 · f72e503 · f72e503
2 parents a8ec2fd + dee13f3
commit f72e503
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/ncolony/process_events.py b/ncolony/process_events.py
@@ -7,6 +7,7 @@
 """
 
 import json
+import os
 
 import six
 
@@ -16,7 +17,7 @@
 
 from ncolony import interfaces
 
-VALID_KEYS = frozenset(['args', 'uid', 'gid', 'env'])
+VALID_KEYS = frozenset(['args', 'uid', 'gid', 'env', 'env_inherit'])
 
 @interface.implementer(interfaces.IMonitorEventReceiver)
 class Receiver(object):
@@ -26,8 +27,11 @@ class Receiver(object):
     :params monitor: a ProcessMonitor
     """
 
-    def __init__(self, monitor):
+    def __init__(self, monitor, environ=None):
         """Initialize from ProcessMonitor"""
+        if environ is None:
+            environ = os.environ
+        self.environ = environ
         self.monitor = monitor
 
     def add(self, name, contents):
@@ -44,6 +48,8 @@ def add(self, name, contents):
                           if key in VALID_KEYS}
         parsedContents['name'] = name
         parsedContents['env'] = parsedContents.get('env', {})
+        for key in parsedContents.pop('env_inherit', []):
+            parsedContents['env'][key] = self.environ.get(key, '')
         parsedContents['env']['NCOLONY_CONFIG'] = contents
         parsedContents['env']['NCOLONY_NAME'] = name
         self.monitor.addProcess(**parsedContents)

diff --git a/ncolony/tests/test_process_events.py b/ncolony/tests/test_process_events.py
@@ -128,6 +128,30 @@ def test_add_with_junk(self):
         self.assertEquals(env, {})
         self.assertEquals(self.logMessages, ['Added monitored process: hello'])
 
+    def test_add_with_inherited_env(self):
+        """Test a process addition with all the optional arguments"""
+        small_environment = dict(PATH='123', PYTHONPATH='456')
+        receiver = process_events.Receiver(self.monitor, small_environment)
+        message = helper.dumps2utf8(dict(args=['/bin/echo', 'hello'],
+                                         env_inherit=['PATH']))
+        receiver.add('hello', message)
+        self.assertEquals(len(self.monitor.events), 1)
+        (tp, name, args, uid, gid, env), = self.monitor.events
+        self.assertEquals(tp, 'ADD')
+        self.assertEquals(name, 'hello')
+        self.assertEquals(args, ['/bin/echo', 'hello'])
+        self.assertEquals(uid, None)
+        self.assertEquals(gid, None)
+        self.assertIn('NCOLONY_CONFIG', env)
+        self.assertEquals(env['NCOLONY_CONFIG'], message)
+        self.assertEquals(self.logMessages, ['Added monitored process: hello'])
+        env.pop('NCOLONY_CONFIG')
+        env.pop('NCOLONY_NAME')
+        sent_environment = small_environment.copy()
+        sent_environment.pop('PYTHONPATH')
+        self.assertEquals(env, sent_environment)
+        self.assertEquals(self.logMessages, ['Added monitored process: hello'])
+
     def test_remove(self):
         """Test a process removal"""
         self.receiver.remove('hello')