test: Add retries to tests for engine flakiness

src/cmd_alias.rs

@@ -432,6 +432,7 @@ mod test {
                 io,
                 debug: false,
                 override_host: None,
+                kcl_retry_config: None,
             };
 
             let cmd_alias = crate::cmd_alias::CmdAlias { subcmd: t.cmd };

src/cmd_auth.rs

@@ -614,6 +614,7 @@ mod test {
                 io,
                 debug: false,
                 override_host: None,
+                kcl_retry_config: None,
             };
             if let Some(h) = &t.host {
                 ctx.override_host = Some(h.clone());

src/cmd_completion.rs

@@ -139,6 +139,7 @@ mod test {
                 io,
                 debug: false,
                 override_host: None,
+                kcl_retry_config: None,
             };
 
             cmd.run(&mut ctx).await.unwrap();

src/cmd_config.rs

@@ -265,6 +265,7 @@ mod test {
                 io,
                 debug: false,
                 override_host: None,
+                kcl_retry_config: None,
             };
 
             let cmd_config = crate::cmd_config::CmdConfig { subcmd: t.cmd };
@@ -301,6 +302,7 @@ mod test {
             io,
             debug: false,
             override_host: None,
+            kcl_retry_config: None,
         };
 
         let mut cmd_config = crate::cmd_config::CmdConfig {

src/cmd_file.rs

@@ -868,6 +868,7 @@ mod test {
                 io,
                 debug: false,
                 override_host: None,
+                kcl_retry_config: None,
             };
 
             let cmd_file = crate::cmd_file::CmdFile { subcmd: t.cmd };

src/cmd_generate.rs

@@ -163,6 +163,7 @@ mod test {
             io,
             debug: false,
             override_host: None,
+            kcl_retry_config: None,
         };
 
         let cmd = crate::cmd_generate::CmdGenerateMarkdown { dir: "".to_string() };
@@ -192,6 +193,7 @@ mod test {
             io,
             debug: false,
             override_host: None,
+            kcl_retry_config: None,
         };
 
         let cmd = crate::cmd_generate::CmdGenerateMarkdown { dir: "".to_string() };

src/cmd_kcl.rs

@@ -159,7 +159,7 @@ impl crate::cmd::Command for CmdKclExport {
         let units: UnitLength = meta_settings.default_length_units.to_kcmc();
         let output_format = get_output_format(&self.output_format, units, self.deterministic);
 
-        if crate::context::Context::use_server_kcl_execution() {
+        let (files, session_data) = if crate::context::Context::use_server_kcl_execution() {
             let (mut responses, session_data) = ctx
                 .run_server_kcl_then_modeling_cmds(
                     "",
@@ -176,37 +176,18 @@ impl crate::cmd::Command for CmdKclExport {
                 anyhow::bail!("Expected Export response from engine");
             };
 
-            for file in files {
-                let path = self.output_dir.join(file.name);
-                std::fs::write(&path, file.contents)?;
-
-                writeln!(ctx.io.out, "Wrote file: {}", path.display())?;
-            }
-
-            if self.show_trace {
-                print_trace_link(&mut ctx.io, &session_data)
-            }
-
-            return Ok(());
-        }
-
-        let client = ctx.api_client("")?;
-        let ectx = kcl_lib::ExecutorContext::new(&client, settings).await?;
-        let mut state = kcl_lib::ExecState::new(&ectx);
-        let session_data = ectx
-            .run(&program, &mut state)
-            .await
-            .map_err(|err| kcl_error_fmt::into_miette(err, &code))?
-            .1;
-        kcl_error_fmt::check_exec_state_issues(
-            &mut ctx.io.err_out,
-            &filepath.display().to_string(),
-            &code,
-            &state,
-            self.run_options.issue_check(),
-        )?;
-
-        let files = ectx.export(output_format).await?;
+            (files, session_data)
+        } else {
+            ctx.run_kcl_then_export(
+                &filepath.display().to_string(),
+                &code,
+                &program,
+                settings,
+                self.run_options.issue_check(),
+                output_format,
+            )
+            .await?
+        };
 
         // Save the files to our export directory.
         for file in files {

src/cmd_say.rs

@@ -116,6 +116,7 @@ mod test {
                 io,
                 debug: false,
                 override_host: None,
+                kcl_retry_config: None,
             };
 
             let cmd_say = crate::cmd_say::CmdSay { input: t.cmd.input };

src/cmd_user.rs

@@ -85,6 +85,7 @@ mod test {
                 io,
                 debug: false,
                 override_host: None,
+                kcl_retry_config: None,
             };
 
             let cmd_user = crate::cmd_user::CmdUser { subcmd: t.cmd };