properly flush batch (#5565)

* properly flush batch

Signed-off-by: Jess Frazelle <[email protected]>

* add test

Signed-off-by: Jess Frazelle <[email protected]>

* add test

Signed-off-by: Jess Frazelle <[email protected]>

* A snapshot a day keeps the bugs away! 📷🐛 (OS: namespace-profile-ubuntu-8-cores)

* only flush if not on error

Signed-off-by: Jess Frazelle <[email protected]>

* better method

Signed-off-by: Jess Frazelle <[email protected]>

* updates

Signed-off-by: Jess Frazelle <[email protected]>

* ensure nothing in batch assembly test

Signed-off-by: Jess Frazelle <[email protected]>

---------

Signed-off-by: Jess Frazelle <[email protected]>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: jessfraz

src/wasm-lib/kcl/src/execution/exec_ast.rs

@@ -123,7 +123,11 @@ impl ExecutorContext {
         };
         std::mem::swap(&mut exec_state.mod_local, &mut local_state);
 
-        if !exec_kind.is_isolated() && new_units != old_units {
+        // We only need to reset the units if we are not on the Main path.
+        // If we reset at the end of the main path, then we just add on an extra
+        // command and we'd need to flush the batch again.
+        // This avoids that.
+        if !exec_kind.is_isolated() && new_units != old_units && *path != ModulePath::Main {
             self.engine.set_units(old_units.into(), Default::default()).await?;
         }
         self.engine.replace_execution_kind(original_execution).await;

src/wasm-lib/kcl/src/simulation_tests.rs

@@ -2074,3 +2074,25 @@ mod import_file_parse_error {
         super::execute(TEST_NAME, true).await
     }
 }
+
+mod flush_batch_on_end {
+    const TEST_NAME: &str = "flush_batch_on_end";
+
+    /// Test parsing KCL.
+    #[test]
+    fn parse() {
+        super::parse(TEST_NAME);
+    }
+
+    /// Test that parsing and unparsing KCL produces the original KCL input.
+    #[test]
+    fn unparse() {
+        super::unparse(TEST_NAME)
+    }
+
+    /// Test that KCL is executed correctly.
+    #[tokio::test(flavor = "multi_thread")]
+    async fn kcl_test_execute() {
+        super::execute(TEST_NAME, true).await
+    }
+}

src/wasm-lib/kcl/tests/cube_with_error/artifact_commands.snap

@@ -1,7 +1,6 @@
 ---
 source: kcl/src/simulation_tests.rs
 description: Artifact commands cube_with_error.kcl
-snapshot_kind: text
 ---
 [
   {