2 files changed, 91 insertions, 4 deletions

diff --git a/lldb/source/Interpreter/Args.cpp b/lldb/source/Interpreter/Args.cpp
index 36ea78cee1d..25bcac81779 100644
--- a/lldb/source/Interpreter/Args.cpp
+++ b/lldb/source/Interpreter/Args.cpp
@@ -976,13 +976,15 @@ void Args::LongestCommonPrefix(std::string &common_prefix) {
 
 void Args::AddOrReplaceEnvironmentVariable(llvm::StringRef env_var_name,
                                            llvm::StringRef new_value) {
-  if (env_var_name.empty() || new_value.empty())
+  if (env_var_name.empty())
     return;
 
   // Build the new entry.
   std::string var_string(env_var_name);
-  var_string += "=";
-  var_string += new_value;
+  if (!new_value.empty()) {
+    var_string += "=";
+    var_string += new_value;
+  }
 
   size_t index = 0;
   if (ContainsEnvironmentVariable(env_var_name, &index)) {
@@ -1006,7 +1008,7 @@ bool Args::ContainsEnvironmentVariable(llvm::StringRef env_var_name,
 
     llvm::StringRef name, value;
     std::tie(name, value) = arg_value.split('=');
-    if (name == env_var_name && !value.empty()) {
+    if (name == env_var_name) {
       if (argument_index)
         *argument_index = i;
       return true;
diff --git a/lldb/unittests/Interpreter/TestArgs.cpp b/lldb/unittests/Interpreter/TestArgs.cpp
index 1dbcad15573..3d1153efff5 100644
--- a/lldb/unittests/Interpreter/TestArgs.cpp
+++ b/lldb/unittests/Interpreter/TestArgs.cpp
@@ -11,6 +11,9 @@
 
 #include "lldb/Interpreter/Args.h"
 
+#include <limits>
+#include <sstream>
+
 using namespace lldb_private;
 
 TEST(ArgsTest, TestSingleArg) {
@@ -153,3 +156,85 @@ TEST(ArgsTest, StringToScriptLanguage) {
 }
 
 TEST(ArgsTest, StringToVersion) {}
+
+// Environment Variable Tests
+
+class EnvVarFixture: public ::testing::Test {
+protected:
+
+    void SetUp() {
+        args.AppendArgument(llvm::StringRef("Arg1=foo"));
+        args.AppendArgument(llvm::StringRef("Arg2"));
+        args.AppendArgument(llvm::StringRef("Arg3=bar"));
+    }
+
+    size_t GetIndexForEnvVar(llvm::StringRef envvar_name) {
+        size_t argument_index = std::numeric_limits<size_t>::max();
+        EXPECT_TRUE(args.ContainsEnvironmentVariable(envvar_name,
+                                                     &argument_index));
+        EXPECT_LT(argument_index, args.GetArgumentCount());
+        return argument_index;
+    }
+
+    Args  args;
+};
+
+
+TEST_F(EnvVarFixture, TestContainsEnvironmentVariableNoValue) {
+    EXPECT_TRUE(args.ContainsEnvironmentVariable(llvm::StringRef("Arg2")));
+}
+
+TEST_F(EnvVarFixture, TestContainsEnvironmentVariableWithValue) {
+    EXPECT_TRUE(args.ContainsEnvironmentVariable(llvm::StringRef("Arg3")));
+}
+
+TEST_F(EnvVarFixture, TestContainsEnvironmentVariableNonExistentVariable) {
+    auto nonexistent_envvar = llvm::StringRef("ThisEnvVarShouldNotExist");
+    EXPECT_FALSE(args.ContainsEnvironmentVariable(nonexistent_envvar));
+}
+
+TEST_F(EnvVarFixture, TestReplaceEnvironmentVariableInitialNoValueWithNoValue) {
+    auto envvar_name = llvm::StringRef("Arg2");
+    auto argument_index = GetIndexForEnvVar(envvar_name);
+
+    args.AddOrReplaceEnvironmentVariable(envvar_name, llvm::StringRef(""));
+    EXPECT_TRUE(args.ContainsEnvironmentVariable(envvar_name));
+    EXPECT_EQ(envvar_name, args.GetArgumentAtIndex(argument_index));
+}
+
+TEST_F(EnvVarFixture, TestReplaceEnvironmentVariableInitialNoValueWithValue) {
+    auto envvar_name = llvm::StringRef("Arg2");
+    auto argument_index = GetIndexForEnvVar(envvar_name);
+
+    auto new_value = llvm::StringRef("NewValue");
+    args.AddOrReplaceEnvironmentVariable(envvar_name, new_value);
+    EXPECT_TRUE(args.ContainsEnvironmentVariable(envvar_name));
+
+    std::stringstream stream;
+    stream << envvar_name.str() << '=' << new_value.str();
+    EXPECT_EQ(llvm::StringRef(stream.str()),
+              args.GetArgumentAtIndex(argument_index));
+}
+
+TEST_F(EnvVarFixture, TestReplaceEnvironmentVariableInitialValueWithNoValue) {
+    auto envvar_name = llvm::StringRef("Arg1");
+    auto argument_index = GetIndexForEnvVar(envvar_name);
+
+    args.AddOrReplaceEnvironmentVariable(envvar_name, llvm::StringRef(""));
+    EXPECT_TRUE(args.ContainsEnvironmentVariable(envvar_name));
+    EXPECT_EQ(envvar_name, args.GetArgumentAtIndex(argument_index));
+}
+
+TEST_F(EnvVarFixture, TestReplaceEnvironmentVariableInitialValueWithValue) {
+    auto envvar_name = llvm::StringRef("Arg1");
+    auto argument_index = GetIndexForEnvVar(envvar_name);
+
+    auto new_value = llvm::StringRef("NewValue");
+    args.AddOrReplaceEnvironmentVariable(envvar_name, new_value);
+    EXPECT_TRUE(args.ContainsEnvironmentVariable(envvar_name));
+
+    std::stringstream stream;
+    stream << envvar_name.str() << '=' << new_value.str();
+    EXPECT_EQ(llvm::StringRef(stream.str()),
+              args.GetArgumentAtIndex(argument_index));
+}