1 files changed, 42 insertions, 6 deletions

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