feat(lib/daemon.py): add way to specify env vars

3 files changed, 126 insertions, 6 deletions

diff --git a/share/python/lib/daemon.md b/share/python/lib/daemon.md
index c4f5e27..c7bb49f 100644
--- a/share/python/lib/daemon.md
+++ b/share/python/lib/daemon.md
@@ -111,6 +111,71 @@ echo hola 1>&2
 echo hi there
 ~~~
 
+# Can specify additional environment variables for daemon
+
+Some daemons are configured through their environment rather than configuration
+files. This scenario verifies that a step can set arbitrary variables in the
+daemon's environment.
+
+~~~scenario
+when I start "/usr/bin/env" as a background process as env, with environment {"custom_variable": "has a Value"}
+when daemon env has produced output
+when I stop background process env
+then daemon env stdout contains "custom_variable=has a Value"
+~~~
+
+~~~scenario
+given a daemon helper shell script env-with-port.py
+when I try to start "./env-with-port.py 8765" as env-with-port, on port 8765, with environment {"custom_variable": "1337"}
+when I stop background process env-with-port
+then daemon env-with-port stdout contains "custom_variable=1337"
+~~~
+
+~~~scenario
+given a daemon helper shell script env-with-port.py
+when I start "./env-with-port.py 8766" as a background process as another-env-with-port, on port 8766, with environment {"subplot2": "000"}
+when daemon another-env-with-port has produced output
+when I stop background process another-env-with-port
+then daemon another-env-with-port stdout contains "subplot2=000"
+~~~
+
+It's important that these new environment variables are not inherited by the
+steps that follow. To verify that, we run one more scenario which *doesn't* set
+any variables, but checks that none of the variables we mentioned above are
+present.
+
+~~~scenario
+when I start "/usr/bin/env" as a background process as env2
+when daemon env2 has produced output
+when I stop background process env2
+then daemon env2 stdout doesn't contain "custom_variable=has a Value"
+then daemon env2 stdout doesn't contain "custom_variable=1337"
+then daemon env2 stdout doesn't contain "subplot2=000"
+~~~
+
+~~~{#env-with-port.py .file .python .numberLines}
+#!/usr/bin/env python3
+
+import os
+import socket
+import sys
+import time
+
+for (key, value) in os.environ.items():
+    print(f"{key}={value}")
+
+port = int(sys.argv[1])
+print(f"port is {port}")
+
+s = socket.socket()
+s.bind(("127.0.0.1", port))
+s.listen()
+
+(conn, _) = s.accept()
+conn.recv(1)
+s.close()
+~~~
+
 
 ---
 title: Acceptance criteria for the lib/daemon Subplot library
diff --git a/share/python/lib/daemon.py b/share/python/lib/daemon.py
index 0573a64..b6658ae 100644
--- a/share/python/lib/daemon.py
+++ b/share/python/lib/daemon.py
@@ -1,3 +1,4 @@
+import json
 import logging
 import os
 import signal
@@ -16,15 +17,15 @@ def _daemon_shell_script(ctx, filename=None):
 
 
 # Start a daemon that will open a port on localhost.
-def daemon_start_on_port(ctx, path=None, args=None, name=None, port=None):
-    _daemon_start(ctx, path=path, args=args, name=name)
+def daemon_start_on_port(ctx, path=None, args=None, name=None, port=None, env=None):
+    _daemon_start(ctx, path=path, args=args, name=name, env=env)
     daemon_wait_for_port("localhost", port)
 
 
 # Start a daemon after a little wait. This is used only for testing the
 # port-waiting code.
-def _daemon_start_soonish(ctx, path=None, args=None, name=None, port=None):
-    _daemon_start(ctx, path=os.path.abspath(path), args=args, name=name)
+def _daemon_start_soonish(ctx, path=None, args=None, name=None, port=None, env=None):
+    _daemon_start(ctx, path=os.path.abspath(path), args=args, name=name, env=env)
     daemon = ctx.declare("_daemon")
 
     # Store the PID of the process we just started so that _daemon_stop_soonish
@@ -41,7 +42,7 @@ def _daemon_start_soonish(ctx, path=None, args=None, name=None, port=None):
     logging.info("pgrep: %r", _daemon_pgrep(path))
 
 
-def _daemon_stop_soonish(ctx, path=None, args=None, name=None, port=None):
+def _daemon_stop_soonish(ctx, path=None, args=None, name=None, port=None, env=None):
     ns = ctx.declare("_daemon")
     pid = ns["_soonish"]
     logging.debug(f"Stopping soonishly-started daemon, {pid}")
@@ -55,7 +56,7 @@ def _daemon_stop_soonish(ctx, path=None, args=None, name=None, port=None):
 # Start a daeamon, get its PID. Don't wait for a port or anything. This is
 # meant for background processes that don't have port. Useful for testing the
 # lib/daemon library of Subplot, but not much else.
-def _daemon_start(ctx, path=None, args=None, name=None):
+def _daemon_start(ctx, path=None, args=None, name=None, env=None):
     runcmd_run = globals()["runcmd_run"]
     runcmd_exit_code_is = globals()["runcmd_exit_code_is"]
     runcmd_get_exit_code = globals()["runcmd_get_exit_code"]
@@ -65,12 +66,19 @@ def _daemon_start(ctx, path=None, args=None, name=None):
     path = os.path.abspath(path)
     argv = [path] + args.split()
 
+    env = json.loads(env or "{}")
+    env_vars = []
+    if env:
+        for (key, value) in env.items():
+            env_vars.extend(["-E", f"{key}={value}"])
+
     logging.debug(f"Starting daemon {name}")
     logging.debug(f"  ctx={ctx.as_dict()}")
     logging.debug(f"  name={name}")
     logging.debug(f"  path={path}")
     logging.debug(f"  args={args}")
     logging.debug(f"  argv={argv}")
+    logging.debug(f"  env={env}")
 
     ns = ctx.declare("_daemon")
 
@@ -96,6 +104,7 @@ def _daemon_start(ctx, path=None, args=None, name=None):
             "-o",
             this["stdout"],
         ]
+        + env_vars
         + argv,
     )
 
@@ -264,6 +273,20 @@ def daemon_stdout_is(ctx, name=None, text=None):
     _daemon_output_is(ctx, name, text, daemon_get_stdout)
 
 
+def daemon_stdout_contains(ctx, name=None, text=None):
+    daemon_get_stdout = globals()["daemon_get_stdout"]
+    assert_eq = globals()["assert_eq"]
+    result = _daemon_output_contains(ctx, name, text, daemon_get_stdout)
+    assert_eq(result, True)
+
+
+def daemon_stdout_doesnt_contain(ctx, name=None, text=None):
+    daemon_get_stdout = globals()["daemon_get_stdout"]
+    assert_eq = globals()["assert_eq"]
+    result = _daemon_output_contains(ctx, name, text, daemon_get_stdout)
+    assert_eq(result, False)
+
+
 def daemon_stderr_is(ctx, name=None, text=None):
     daemon_get_stderr = globals()["daemon_get_stderr"]
     _daemon_output_is(ctx, name, text, daemon_get_stderr)
@@ -274,3 +297,16 @@ def _daemon_output_is(ctx, name, text, getter):
     text = bytes(text, "UTF-8").decode("unicode_escape")
     output = getter(ctx, name)
     assert_eq(text, output)
+
+
+def _daemon_output_contains(ctx, name, text, getter):
+    text = bytes(text, "UTF-8").decode("unicode_escape")
+    output = getter(ctx, name)
+
+    logging.debug("  output:")
+    log_lines(indent, output)
+
+    logging.debug("  text:")
+    log_lines(indent, text)
+
+    return text in output
diff --git a/share/python/lib/daemon.yaml b/share/python/lib/daemon.yaml
index 424a98f..f4a2f86 100644
--- a/share/python/lib/daemon.yaml
+++ b/share/python/lib/daemon.yaml
@@ -7,13 +7,26 @@
 - when: I start "{path}{args:text}" as a background process as {name}, on port {port}
   function: daemon_start_on_port
 
+- when: I start "(?P<path>[^ "]+)(?P<args>[^"]*)" as a background process as (?P<name>[^,]+), on port (?P<port>\d+), with environment (?P<env>.*)
+  regex: true
+  function: daemon_start_on_port
+
 - when: I try to start "{path}{args:text}" as {name}, on port {port}
   function: _daemon_start_soonish
   cleanup: _daemon_stop_soonish
 
+- when: I try to start "(?P<path>[^ "]+)(?P<args>[^"]*)" as (?P<name>[^,]+), on port (?P<port>\d+), with environment (?P<env>.*)
+  regex: true
+  function: _daemon_start_soonish
+  cleanup: _daemon_stop_soonish
+
 - when: I start "{path}{args:text}" as a background process as {name}
   function: _daemon_start
 
+- when: I start "(?P<path>[^ "]+)(?P<args>[^"]*)" as a background process as (?P<name>[^,]+), with environment (?P<env>.*)
+  regex: true
+  function: _daemon_start
+
 - when: I stop background process {name}
   function: daemon_stop
 
@@ -35,6 +48,12 @@
 - then: daemon {name} stdout is "{text:text}"
   function: daemon_stdout_is
 
+- then: daemon {name} stdout contains "{text:text}"
+  function: daemon_stdout_contains
+
+- then: daemon {name} stdout doesn't contain "{text:text}"
+  function: daemon_stdout_doesnt_contain
+
 - then: daemon {name} stderr is "{text:text}"
   function: daemon_stderr_is