Refactor Env and add unit tests (#2010)

mobile-dev-inc · Sep 3, 2024 · c95fe9d · c95fe9d
1 parent 2c3efe7
commit c95fe9d
Show file tree

Hide file tree

Showing 9 changed files with 98 additions and 33 deletions.
diff --git a/maestro-cli/src/main/java/maestro/cli/command/CloudCommand.kt b/maestro-cli/src/main/java/maestro/cli/command/CloudCommand.kt
@@ -33,6 +33,7 @@ import picocli.CommandLine.Option
 import java.io.File
 import java.util.concurrent.Callable
 import java.util.concurrent.TimeUnit
+import maestro.orchestra.util.Env.withDefaultEnvVars
 
 @CommandLine.Command(
     name = "cloud",
@@ -172,6 +173,10 @@ class CloudCommand : Callable<Int> {
             }
         }
 
+        env = env
+            .withInjectedShellEnvVars()
+            .withDefaultEnvVars(flowsFile)
+
         return CloudInteractor(
             client = ApiClient(apiUrl),
             failOnTimeout = failOnTimeout,
@@ -181,7 +186,7 @@ class CloudCommand : Callable<Int> {
             flowFile = flowsFile,
             appFile = appFile,
             mapping = mapping,
-            env = env.withInjectedShellEnvVars(),
+            env = env,
             uploadName = uploadName,
             repoOwner = repoOwner,
             repoName = repoName,

diff --git a/maestro-cli/src/main/java/maestro/cli/command/TestCommand.kt b/maestro-cli/src/main/java/maestro/cli/command/TestCommand.kt
@@ -59,6 +59,7 @@ import java.util.concurrent.CountDownLatch
 import kotlin.io.path.absolutePathString
 import kotlin.system.exitProcess
 import kotlin.time.Duration.Companion.seconds
+import maestro.orchestra.util.Env.withDefaultEnvVars
 
 @CommandLine.Command(
     name = "test",
@@ -161,7 +162,9 @@ class TestCommand : Callable<Int> {
             throw CliError(e.message)
         }
 
-        env = env.withInjectedShellEnvVars()
+        env = env
+            .withInjectedShellEnvVars()
+            .withDefaultEnvVars(flowFile)
 
         TestDebugReporter.install(
             debugOutputPathAsString = debugOutput,

diff --git a/maestro-cli/src/main/java/maestro/cli/command/UploadCommand.kt b/maestro-cli/src/main/java/maestro/cli/command/UploadCommand.kt
@@ -29,6 +29,7 @@ import picocli.CommandLine
 import picocli.CommandLine.Option
 import java.io.File
 import java.util.concurrent.Callable
+import maestro.orchestra.util.Env.withDefaultEnvVars
 
 @CommandLine.Command(
     name = "upload",
@@ -89,14 +90,18 @@ class UploadCommand : Callable<Int> {
                 .fgDefault()
         )
 
+        env = env
+            .withInjectedShellEnvVars()
+            .withDefaultEnvVars(flowFile)
+
         return CloudInteractor(
             client = ApiClient(apiUrl),
         ).upload(
             async = true,
             flowFile = flowFile,
             appFile = appFile,
             mapping = mapping,
-            env = env.withInjectedShellEnvVars(),
+            env = env,
             uploadName = uploadName,
             repoOwner = repoOwner,
             repoName = repoName,

diff --git a/maestro-cli/src/main/java/maestro/cli/runner/TestRunner.kt b/maestro-cli/src/main/java/maestro/cli/runner/TestRunner.kt
@@ -23,7 +23,6 @@ import org.slf4j.LoggerFactory
 import java.io.File
 import java.nio.file.Path
 import kotlin.concurrent.thread
-import maestro.orchestra.util.Env.withDefaultEnvVars
 
 /**
  * Knows how to run a single Maestro flow (either one-shot or continuously).
@@ -53,7 +52,7 @@ object TestRunner {
 
         val result = runCatching(resultView, maestro) {
             val commands = YamlCommandReader.readCommands(flowFile.toPath())
-                .withEnv(env.withDefaultEnvVars(flowFile.nameWithoutExtension))
+                .withEnv(env)
 
             YamlCommandReader.getConfig(commands)?.name?.let {
                 aiOutput = aiOutput.copy(flowName = it)
@@ -109,7 +108,7 @@ object TestRunner {
 
                 val commands = YamlCommandReader
                     .readCommands(flowFile.toPath())
-                    .withEnv(env.withDefaultEnvVars(flowFile.nameWithoutExtension))
+                    .withEnv(env)
 
                 // Restart the flow if anything has changed
                 if (commands != previousCommands) {

diff --git a/maestro-cli/src/main/java/maestro/cli/runner/TestSuiteInteractor.kt b/maestro-cli/src/main/java/maestro/cli/runner/TestSuiteInteractor.kt
@@ -28,7 +28,6 @@ import java.io.File
 import java.nio.file.Path
 import kotlin.system.measureTimeMillis
 import kotlin.time.Duration.Companion.seconds
-import maestro.orchestra.util.Env.withDefaultEnvVars
 
 /**
  * Similar to [TestRunner], but:
@@ -192,7 +191,7 @@ class TestSuiteInteractor(
             try {
                 val commands = YamlCommandReader
                     .readCommands(flowFile.toPath())
-                    .withEnv(env.withDefaultEnvVars(flowFile.nameWithoutExtension))
+                    .withEnv(env)
 
                 YamlCommandReader.getConfig(commands)?.name?.let { flowName = it }
 

diff --git a/maestro-orchestra-models/src/main/java/maestro/orchestra/util/Env.kt b/maestro-orchestra-models/src/main/java/maestro/orchestra/util/Env.kt
@@ -1,5 +1,6 @@
 package maestro.orchestra.util
 
+import java.io.File
 import maestro.js.JsEngine
 import maestro.orchestra.DefineVariablesCommand
 import maestro.orchestra.MaestroCommand
@@ -24,27 +25,19 @@ object Env {
             }
     }
 
-    fun List<MaestroCommand>.withEnv(env: Map<String, String>): List<MaestroCommand> {
-        if (env.isEmpty()) {
-            return this
-        }
+    fun List<MaestroCommand>.withEnv(env: Map<String, String>): List<MaestroCommand> =
+        if (env.isEmpty()) this
+        else listOf(MaestroCommand(DefineVariablesCommand(env))) + this
 
-        return listOf(MaestroCommand(DefineVariablesCommand(env))) + this
-    }
-
-    fun Map<String, String>.withInjectedShellEnvVars(): Map<String, String> {
-        val mutable = this.toMutableMap()
-        val sysEnv = System.getenv()
-        for (sysEnvKey in sysEnv.keys.filter { it.startsWith("MAESTRO_") }) {
-            val sysEnvValue = sysEnv[sysEnvKey]
-            if (mutable[sysEnvKey] == null && sysEnvValue != null) {
-                mutable[sysEnvKey] = sysEnvValue
-            }
-        }
+    fun Map<String, String>.withInjectedShellEnvVars(): Map<String, String> = this +
+        System.getenv()
+            .filterKeys { it.startsWith("MAESTRO_") && this.containsKey(it).not() }
+            .filterValues { it != null && it.isNotEmpty() }
 
-        return mutable
+    fun Map<String, String>.withDefaultEnvVars(flowFile: File? = null): Map<String, String> {
+        val defaultEnvVars = mutableMapOf<String, String>()
+        flowFile?.nameWithoutExtension?.let { defaultEnvVars["MAESTRO_FILENAME"] = it }
+        return if (defaultEnvVars.isEmpty()) this
+        else this + defaultEnvVars
     }
-
-    fun Map<String, String>.withDefaultEnvVars(fileName: String? = null) =
-        fileName?.let { this + mapOf("MAESTRO_FILENAME" to fileName) } ?: this
 }
diff --git a/maestro-orchestra-models/src/test/kotlin/maestro/orchestra/util/EnvTest.kt b/maestro-orchestra-models/src/test/kotlin/maestro/orchestra/util/EnvTest.kt
@@ -0,0 +1,63 @@
+package maestro.orchestra.util
+
+import com.google.common.truth.Truth.assertThat
+import java.io.File
+import kotlin.random.Random
+import maestro.orchestra.ApplyConfigurationCommand
+import maestro.orchestra.DefineVariablesCommand
+import maestro.orchestra.MaestroCommand
+import maestro.orchestra.MaestroConfig
+import maestro.orchestra.util.Env.withDefaultEnvVars
+import maestro.orchestra.util.Env.withEnv
+import maestro.orchestra.util.Env.withInjectedShellEnvVars
+import org.junit.jupiter.api.Test
+
+class EnvTest {
+
+    private val emptyEnv = emptyMap<String, String>()
+
+    @Test
+    fun `withDefaultEnvVars should add file name without extension`() {
+        val env = emptyEnv.withDefaultEnvVars(File("myFlow.yml"))
+        assertThat(env["MAESTRO_FILENAME"]).isEqualTo("myFlow")
+    }
+
+    @Test
+    fun `withDefaultEnvVars should override MAESTRO_FILENAME`() {
+        val env = mapOf("MAESTRO_FILENAME" to "otherFile").withDefaultEnvVars(File("myFlow.yml"))
+        assertThat(env["MAESTRO_FILENAME"]).isEqualTo("myFlow")
+    }
+
+    @Test
+    fun `withInjectedShellEnvVars only keeps MAESTRO_ vars`() {
+        val env = emptyEnv.withInjectedShellEnvVars()
+        assertThat(env.filterKeys { it.startsWith("MAESTRO_").not() }).isEmpty()
+    }
+
+    @Test
+    fun `withInjectedShellEnvVars does not strip previous MAESTRO_ vars`() {
+        val rand = Random.nextInt()
+        val env = mapOf("MAESTRO_$rand" to "$rand").withInjectedShellEnvVars()
+        assertThat(env["MAESTRO_$rand"]).isEqualTo("$rand")
+    }
+
+    @Test
+    fun `withEnv does not affect empty env`() {
+        val commands = emptyList<MaestroCommand>()
+
+        val withEnv = commands.withEnv(emptyEnv)
+
+        assertThat(withEnv).isEmpty()
+    }
+
+    @Test
+    fun `withEnv prepends DefineVariable command`() {
+        val env = mapOf("MY_ENV_VAR" to "1234")
+        val applyConfig = MaestroCommand(ApplyConfigurationCommand(MaestroConfig()))
+        val defineVariables = MaestroCommand(DefineVariablesCommand(env))
+
+        val withEnv = listOf(applyConfig).withEnv(env)
+
+        assertThat(withEnv).containsExactly(defineVariables, applyConfig)
+    }
+}
diff --git a/maestro-orchestra/src/main/java/maestro/orchestra/yaml/YamlCommandReader.kt b/maestro-orchestra/src/main/java/maestro/orchestra/yaml/YamlCommandReader.kt
@@ -30,14 +30,12 @@ import java.nio.file.Path
 import kotlin.io.path.absolute
 import kotlin.io.path.absolutePathString
 import kotlin.io.path.isDirectory
-import kotlin.io.path.nameWithoutExtension
 import kotlin.io.path.readText
 import maestro.orchestra.ApplyConfigurationCommand
 import maestro.orchestra.MaestroCommand
 import maestro.orchestra.MaestroConfig
 import maestro.orchestra.WorkspaceConfig
 import maestro.orchestra.error.SyntaxError
-import maestro.orchestra.util.Env.withDefaultEnvVars
 import maestro.orchestra.util.Env.withEnv
 
 object YamlCommandReader {

diff --git a/maestro-test/src/test/kotlin/maestro/test/IntegrationTest.kt b/maestro-test/src/test/kotlin/maestro/test/IntegrationTest.kt
@@ -28,7 +28,6 @@ import org.junit.jupiter.api.assertThrows
 import java.awt.Color
 import java.io.File
 import java.nio.file.Paths
-import maestro.orchestra.error.SyntaxError
 import kotlin.system.measureTimeMillis
 import maestro.orchestra.util.Env.withDefaultEnvVars
 
@@ -3170,8 +3169,9 @@ class IntegrationTest {
     private fun readCommands(caseName: String, withEnv: () -> Map<String, String> = { emptyMap() }): List<MaestroCommand> {
         val resource = javaClass.classLoader.getResource("$caseName.yaml")
             ?: throw IllegalArgumentException("File $caseName.yaml not found")
-        return YamlCommandReader.readCommands(Paths.get(resource.toURI()))
-            .withEnv(withEnv().withDefaultEnvVars(caseName))
+        val flowPath = Paths.get(resource.toURI())
+        return YamlCommandReader.readCommands(flowPath)
+            .withEnv(withEnv().withDefaultEnvVars(flowPath.toFile()))
     }
 
 }