Fix Environment.updating for custom and rawBytes (#199)

Author: a7medev

Sources/Subprocess/Configuration.swift

@@ -422,7 +422,32 @@ public struct Environment: Sendable, Hashable {
     /// Keys with `nil` values in `newValue` will be removed from existing
     /// `Environment` before passing to child process.
     public func updating(_ newValue: [Key: String?]) -> Self {
-        return .init(config: .inherit(newValue))
+        switch config {
+        case .inherit(var overrides):
+            for (key, value) in newValue {
+                overrides[key] = value
+            }
+            return .init(config: .inherit(overrides))
+        case .custom(var environment):
+            for (key, value) in newValue {
+                environment[key] = value
+            }
+            return .init(config: .custom(environment))
+        #if !os(Windows)
+        case .rawBytes(var rawBytesArray):
+            let overriddenKeys = newValue.keys.map { Array("\($0)=".utf8) }
+            rawBytesArray.removeAll {
+                overriddenKeys.contains(where: $0.starts)
+            }
+
+            for (key, value) in newValue {
+                if let value {
+                    rawBytesArray.append(Array("\(key)=\(value)\0".utf8))
+                }
+            }
+            return .init(config: .rawBytes(rawBytesArray))
+        #endif
+        }
     }
     /// Use custom environment variables
     public static func custom(_ newValue: [Key: String]) -> Self {

Tests/SubprocessTests/IntegrationTests.swift

@@ -435,6 +435,132 @@ extension SubprocessIntegrationTests {
         )
         #endif
     }
+
+    @Test func testEnvironmentCustomUpdating() async throws {
+        #if os(Windows)
+        let pathValue = ProcessInfo.processInfo.environment["Path"] ?? ProcessInfo.processInfo.environment["PATH"] ?? ProcessInfo.processInfo.environment["path"]
+        let customPath = "C:\\Custom\\Path"
+        let setup = TestSetup(
+            executable: .name("cmd.exe"),
+            arguments: ["/c", "echo %CUSTOMPATH%"],
+            environment: .custom([
+                "Path": try #require(pathValue),
+                "ComSpec": try #require(ProcessInfo.processInfo.environment["ComSpec"]),
+            ]).updating([
+                "CUSTOMPATH": customPath
+            ])
+        )
+        #else
+        let customPath = "/custom/path"
+        let setup = TestSetup(
+            executable: .path("/bin/sh"),
+            arguments: ["-c", "printenv CUSTOMPATH"],
+            environment: .custom([
+                "PATH": "/bin:/usr/bin"
+            ]).updating([
+                "CUSTOMPATH": customPath
+            ])
+        )
+        #endif
+        let result = try await _run(
+            setup,
+            input: .none,
+            output: .string(limit: 32),
+            error: .discarded
+        )
+        #expect(result.terminationStatus.isSuccess)
+        let output = result.standardOutput?
+            .trimmingNewLineAndQuotes()
+        #expect(
+            output == customPath
+        )
+    }
+
+    @Test func testEnvironmentCustomUpdatingUnsetValue() async throws {
+        #if os(Windows)
+        let pathValue = ProcessInfo.processInfo.environment["Path"] ?? ProcessInfo.processInfo.environment["PATH"] ?? ProcessInfo.processInfo.environment["path"]
+        let setup = TestSetup(
+            executable: .name("cmd.exe"),
+            arguments: ["/c", "echo %REMOVEME%"],
+            environment: .custom([
+                "Path": try #require(pathValue),
+                "ComSpec": try #require(ProcessInfo.processInfo.environment["ComSpec"]),
+                "REMOVEME": "value",
+            ]).updating([
+                "REMOVEME": nil
+            ])
+        )
+        #else
+        let setup = TestSetup(
+            executable: .path("/bin/sh"),
+            arguments: ["-c", "printenv REMOVEME"],
+            environment: .custom([
+                "PATH": "/bin:/usr/bin",
+                "REMOVEME": "value",
+            ]).updating([
+                "REMOVEME": nil
+            ])
+        )
+        #endif
+        let result = try await _run(
+            setup,
+            input: .none,
+            output: .string(limit: 32),
+            error: .discarded
+        )
+        #if os(Windows)
+        #expect(result.standardOutput?.trimmingNewLineAndQuotes() == "%REMOVEME%")
+        #else
+        #expect(result.terminationStatus == .exited(1))
+        #endif
+    }
+
+    #if !os(Windows)
+    @Test func testEnvironmentRawBytesUpdating() async throws {
+        let customValue = "rawbytes_value"
+        let setup = TestSetup(
+            executable: .path("/bin/sh"),
+            arguments: ["-c", "printenv CUSTOMVAR"],
+            environment: .custom([
+                Array("PATH=/bin:/usr/bin\0".utf8)
+            ]).updating([
+                "CUSTOMVAR": customValue
+            ])
+        )
+        let result = try await _run(
+            setup,
+            input: .none,
+            output: .string(limit: 32),
+            error: .discarded
+        )
+        #expect(result.terminationStatus.isSuccess)
+        let output = result.standardOutput?
+            .trimmingNewLineAndQuotes()
+        #expect(
+            output == customValue
+        )
+    }
+
+    @Test func testEnvironmentRawBytesUpdatingUnsetValue() async throws {
+        let setup = TestSetup(
+            executable: .path("/bin/sh"),
+            arguments: ["-c", "printenv REMOVEME"],
+            environment: .custom([
+                Array("PATH=/bin:/usr/bin\0".utf8),
+                Array("REMOVEME=value\0".utf8),
+            ]).updating([
+                "REMOVEME": nil
+            ])
+        )
+        let result = try await _run(
+            setup,
+            input: .none,
+            output: .string(limit: 32),
+            error: .discarded
+        )
+        #expect(result.terminationStatus == .exited(1))
+    }
+    #endif
 }
 
 // MARK: - Working Directory Tests