From 369c2abfbe9848feee82041b3713512fa4b8ff62 Mon Sep 17 00:00:00 2001
From: Michael Carroll <mjcarroll@intrinsic.ai>
Date: Tue, 27 Jun 2023 19:11:18 +0000
Subject: [PATCH] Also check environment when running

Signed-off-by: Michael Carroll <mjcarroll@intrinsic.ai>
---
 include/gz/utils/Subprocess.hh                | 39 +++++++++++++++---
 test/integration/CMakeLists.txt               |  2 +
 .../integration/subprocess/subprocess_main.cc | 10 ++++-
 test/integration/subprocess_TEST.cc           | 41 ++++++++++++++++++-
 4 files changed, 83 insertions(+), 9 deletions(-)

diff --git a/include/gz/utils/Subprocess.hh b/include/gz/utils/Subprocess.hh
index a80c4e4..524dc7f 100644
--- a/include/gz/utils/Subprocess.hh
+++ b/include/gz/utils/Subprocess.hh
@@ -34,10 +34,22 @@ namespace utils
 
 class Subprocess
 {
-  public: Subprocess(std::vector<std::string> _commandLine):
+  public: Subprocess(const std::vector<std::string> &_commandLine,
+                     const std::vector<std::string> &_environment = {}):
     commandLine(_commandLine),
-    process(new subprocess_s)
+    environment(_environment)
   {
+    this->Create();
+
+  }
+
+  private: void Create()
+  {
+    if (this->process)
+        return;
+
+    this->process = new subprocess_s;
+
     std::vector<const char*> commandLineCstr;
     for (const auto &val : this->commandLine)
     {
@@ -45,7 +57,23 @@ class Subprocess
     }
     commandLineCstr.push_back(nullptr);
 
-    auto ret = subprocess_create(commandLineCstr.data(), 0, this->process);
+    std::vector<const char*> environmentCstr;
+    for (const auto &val : this->environment)
+    {
+      environmentCstr.push_back(val.c_str());
+    }
+    environmentCstr.push_back(nullptr);
+
+    int ret = -1;
+    if (this->environment.size())
+    {
+      ret = subprocess_create_ex(commandLineCstr.data(), 0, environmentCstr.data(), this->process);
+    }
+    else
+    {
+      ret = subprocess_create(commandLineCstr.data(), 0, this->process);
+    }
+
     if (0 != ret)
     {
       std::cerr << "failed to create subprocess" << std::endl;
@@ -121,15 +149,14 @@ class Subprocess
         std::cerr << "Failed to join subprocess" << std::endl;
       }
     }
-    else
-    {
-    }
 
     return return_code;
   }
 
   protected: std::vector<std::string> commandLine;
 
+  protected: std::vector<std::string> environment;
+
   protected: subprocess_s * process {nullptr};
 };
 }  // namespace utils
diff --git a/test/integration/CMakeLists.txt b/test/integration/CMakeLists.txt
index db5b686..87a888e 100644
--- a/test/integration/CMakeLists.txt
+++ b/test/integration/CMakeLists.txt
@@ -15,10 +15,12 @@ add_executable(subprocess_executable subprocess/subprocess_main.cc )
 target_include_directories(subprocess_executable PUBLIC
   ${PROJECT_SOURCE_DIR}/cli/include/
   ${PROJECT_SOURCE_DIR}/cli/include/vendored-cli/)
+target_link_libraries(subprocess_executable ${PROJECT_LIBRARY_TARGET_NAME})
 
 if(TARGET INTEGRATION_subprocess_TEST)
   target_compile_definitions(INTEGRATION_subprocess_TEST PRIVATE
     "SUBPROCESS_EXECUTABLE_PATH=\"$<TARGET_FILE:subprocess_executable>\"")
+
 endif()
 
 add_subdirectory(implptr)
diff --git a/test/integration/subprocess/subprocess_main.cc b/test/integration/subprocess/subprocess_main.cc
index e4d14d5..c4c661b 100644
--- a/test/integration/subprocess/subprocess_main.cc
+++ b/test/integration/subprocess/subprocess_main.cc
@@ -19,6 +19,7 @@
 #include <string>
 #include <thread>
 
+#include <gz/utils/Environment.hh>
 #include <gz/utils/cli/CLI.hpp>
 
 class OutputSink
@@ -53,10 +54,10 @@ int main(int argc, char **argv)
   std::string output;
   app.add_option("--output", output, "output destination");
 
-  int iters;
+  int iters = 0;
   app.add_option("--iterations", iters, "number of iterations to run");
 
-  int iter_ms;
+  int iter_ms = 0;
   app.add_option("--iteration-ms", iter_ms, "length of one iteration");
 
   CLI11_PARSE(app, argc, argv);
@@ -68,4 +69,9 @@ int main(int argc, char **argv)
     std::this_thread::sleep_for(std::chrono::milliseconds(iter_ms));
   }
 
+  std::string env_var;
+  if(gz::utils::env("ENV_VAR", env_var))
+  {
+    sink.Write("ENV_VAR=" + env_var);
+  }
 }
diff --git a/test/integration/subprocess_TEST.cc b/test/integration/subprocess_TEST.cc
index 6277820..03641fb 100644
--- a/test/integration/subprocess_TEST.cc
+++ b/test/integration/subprocess_TEST.cc
@@ -39,7 +39,7 @@ TEST(Subprocess, CreateValid)
   auto proc = Subprocess({kExecutablePath, "--help"});
 
   // Sleep for just a bit to guarantee executable finishes
-  std::this_thread::sleep_for(std::chrono::milliseconds(5));
+  std::this_thread::sleep_for(std::chrono::milliseconds(100));
   EXPECT_FALSE(proc.Alive());
 
   auto cout = proc.Stdout();
@@ -92,3 +92,42 @@ TEST(Subprocess, Cerr)
   EXPECT_TRUE(cout.empty());
   EXPECT_FALSE(cerr.empty());
 }
+
+/////////////////////////////////////////////////
+TEST(Subprocess, Environment)
+{
+  {
+    auto proc = Subprocess({kExecutablePath, "--output=cout"},
+                           {"ENV_VAR=foobar"});
+    // Block until the executable is done
+    auto ret = proc.Join();
+    EXPECT_EQ(0u, ret);
+
+    auto cout = proc.Stdout();
+    EXPECT_NE(std::string::npos, cout.find("ENV_VAR=foobar"));
+  }
+
+  {
+    auto proc = Subprocess({kExecutablePath, "--output=cerr"},
+                           {"ENV_VAR=foobar2"});
+    // Block until the executable is done
+    auto ret = proc.Join();
+    EXPECT_EQ(0u, ret);
+
+    auto cerr = proc.Stderr();
+    EXPECT_NE(std::string::npos, cerr.find("ENV_VAR=foobar2"));
+  }
+
+  {
+    auto proc = Subprocess({kExecutablePath},
+                           {"ENV_VAR=foobar3"});
+    // Block until the executable is done
+    auto ret = proc.Join();
+    EXPECT_EQ(0u, ret);
+
+    auto cout = proc.Stdout();
+    auto cerr = proc.Stderr();
+    EXPECT_EQ(std::string::npos, cerr.find("ENV_VAR=foobar3"));
+    EXPECT_EQ(std::string::npos, cout.find("ENV_VAR=foobar3"));
+  }
+}