Skip to main content

kcl_lib/
test_server.rs

1//! Types used to send data to the test server.
2
3use std::path::PathBuf;
4
5use kittycad_modeling_cmds::websocket::RawFile;
6
7use crate::ConnectionError;
8use crate::ExecError;
9use crate::KclError;
10use crate::KclErrorWithOutputs;
11use crate::Program;
12use crate::engine::new_zoo_client;
13use crate::errors::ExecErrorWithState;
14use crate::execution::EnvironmentRef;
15use crate::execution::ExecState;
16use crate::execution::ExecutorContext;
17use crate::execution::ExecutorSettings;
18
19#[derive(serde::Deserialize, serde::Serialize)]
20pub struct RequestBody {
21    pub kcl_program: String,
22    #[serde(default)]
23    pub test_name: String,
24}
25
26/// Executes a kcl program and takes a snapshot of the result.
27/// This returns the bytes of the snapshot.
28pub async fn execute_and_snapshot(code: &str, current_file: Option<PathBuf>) -> Result<image::DynamicImage, ExecError> {
29    let ctx = new_context(true, current_file).await?;
30    let program = Program::parse_no_errs(code).map_err(KclErrorWithOutputs::no_outputs)?;
31    let res = do_execute_and_snapshot(&ctx, program, None)
32        .await
33        .map(|(_, _, snap)| snap)
34        .map_err(|err| err.error);
35    ctx.close().await;
36    res
37}
38
39/// Executes a KCL program. Only returns success or error.
40pub async fn execute(code: &str, current_file: Option<PathBuf>) -> Result<(), ExecError> {
41    let ctx = new_context(true, current_file).await?;
42    let program = Program::parse_no_errs(code).map_err(KclErrorWithOutputs::no_outputs)?;
43    let res = do_execute(&ctx, program, None)
44        .await
45        .map(|_| ())
46        .map_err(|err| err.error);
47    ctx.close().await;
48    res
49}
50
51pub struct Snapshot3d {
52    /// Bytes of the snapshot.
53    pub image: image::DynamicImage,
54    /// Various GLTF files for the resulting export.
55    pub gltf: Vec<RawFile>,
56}
57
58async fn export_gltf_if_requested<F, Fut>(request_gltf: bool, export: F) -> Result<Vec<RawFile>, KclError>
59where
60    F: FnOnce() -> Fut,
61    Fut: std::future::Future<Output = Result<Vec<RawFile>, KclError>>,
62{
63    if request_gltf { export().await } else { Ok(Vec::new()) }
64}
65
66/// Executes a kcl program and takes a snapshot of the result.
67pub async fn execute_and_snapshot_3d(
68    code: &str,
69    current_file: Option<PathBuf>,
70    request_gltf: bool,
71) -> Result<Snapshot3d, ExecError> {
72    let ctx = new_context(true, current_file).await?;
73    let program = Program::parse_no_errs(code).map_err(KclErrorWithOutputs::no_outputs)?;
74    let image = do_execute_and_snapshot(&ctx, program, None)
75        .await
76        .map(|(_, _, snap)| snap)
77        .map_err(|err| err.error)?;
78    let gltf_res = export_gltf_if_requested(request_gltf, || {
79        ctx.export(kittycad_modeling_cmds::format::OutputFormat3d::Gltf(Default::default()))
80    })
81    .await;
82    let gltf = match gltf_res {
83        Err(err) if err.message() == "Nothing to export" => Vec::new(),
84        Err(err) => {
85            eprintln!("Error exporting: {}", err.message());
86            Vec::new()
87        }
88        Ok(x) => x,
89    };
90    ctx.close().await;
91    Ok(Snapshot3d { image, gltf })
92}
93
94/// Executes a kcl program and takes a snapshot of the result.
95/// This returns the bytes of the snapshot.
96#[cfg(test)]
97pub async fn execute_and_snapshot_ast(
98    ast: Program,
99    current_file: Option<PathBuf>,
100    with_export_step: bool,
101    deprecation_version_override: Option<&str>,
102) -> Result<
103    (
104        ExecState,
105        ExecutorContext,
106        EnvironmentRef,
107        image::DynamicImage,
108        Option<Vec<u8>>,
109    ),
110    ExecErrorWithState,
111> {
112    let result = execute_and_snapshot_ast_with_heartbeats(
113        ast,
114        current_file,
115        with_export_step,
116        deprecation_version_override,
117        None,
118    )
119    .await;
120    if let Ok((_, ctx, _, _, _)) = &result {
121        ctx.close().await;
122    }
123    result
124}
125
126/// Executes a KCL program and takes a snapshot without closing the engine
127/// connection. If OK, the caller must close the returned context.
128/// If Err, the context will already be closed within this function.
129#[cfg(test)]
130pub async fn execute_and_snapshot_ast_no_close(
131    ast: Program,
132    current_file: Option<PathBuf>,
133    with_export_step: bool,
134    deprecation_version_override: Option<&str>,
135) -> Result<
136    (
137        ExecState,
138        ExecutorContext,
139        EnvironmentRef,
140        image::DynamicImage,
141        Option<Vec<u8>>,
142    ),
143    ExecErrorWithState,
144> {
145    execute_and_snapshot_ast_with_heartbeats(
146        ast,
147        current_file,
148        with_export_step,
149        deprecation_version_override,
150        Some(5),
151    )
152    .await
153}
154
155#[cfg(test)]
156async fn execute_and_snapshot_ast_with_heartbeats(
157    ast: Program,
158    current_file: Option<PathBuf>,
159    with_export_step: bool,
160    deprecation_version_override: Option<&str>,
161    heartbeats: Option<u64>,
162) -> Result<
163    (
164        ExecState,
165        ExecutorContext,
166        EnvironmentRef,
167        image::DynamicImage,
168        Option<Vec<u8>>,
169    ),
170    ExecErrorWithState,
171> {
172    let ctx = new_context_with_heartbeats(true, current_file, heartbeats).await?;
173    let (exec_state, env, img) = match do_execute_and_snapshot(&ctx, ast, deprecation_version_override).await {
174        Ok((exec_state, env_ref, img)) => (exec_state, env_ref, img),
175        Err(err) => {
176            // If there was an error executing the program, return it.
177            // Close the context to avoid any resource leaks.
178            ctx.close().await;
179            return Err(err);
180        }
181    };
182    let mut step = None;
183    if with_export_step {
184        let files = match ctx.export_step(true).await {
185            Ok(f) => f,
186            Err(err) => {
187                // Close the context to avoid any resource leaks.
188                ctx.close().await;
189                return Err(ExecErrorWithState::new(
190                    ExecError::BadExport(format!("Export failed: {err:?}")),
191                    exec_state.clone(),
192                    None,
193                ));
194            }
195        };
196
197        step = files.into_iter().next().map(|f| f.contents);
198    }
199    Ok((exec_state, ctx, env, img, step))
200}
201
202pub async fn execute_and_snapshot_no_auth(
203    code: &str,
204    current_file: Option<PathBuf>,
205) -> Result<(image::DynamicImage, EnvironmentRef), ExecError> {
206    let ctx = new_context(false, current_file).await?;
207    let program = Program::parse_no_errs(code).map_err(KclErrorWithOutputs::no_outputs)?;
208    let res = do_execute_and_snapshot(&ctx, program, None)
209        .await
210        .map(|(_, env_ref, snap)| (snap, env_ref))
211        .map_err(|err| err.error);
212    ctx.close().await;
213    res
214}
215
216async fn do_execute(
217    ctx: &ExecutorContext,
218    program: Program,
219    _deprecation_version_override: Option<&str>,
220) -> Result<(ExecState, EnvironmentRef), ExecErrorWithState> {
221    let mut exec_state = ExecState::new(ctx);
222    #[cfg(test)]
223    exec_state.set_deprecation_version_override(_deprecation_version_override);
224    let result = ctx.run(&program, &mut exec_state).await;
225    let responses = if result.is_err() {
226        #[cfg(feature = "snapshot-engine-responses")]
227        {
228            Some(exec_state.take_root_module_responses())
229        }
230        #[cfg(not(feature = "snapshot-engine-responses"))]
231        None
232    } else {
233        None
234    };
235    let result = result.map_err(|err| ExecErrorWithState::new(err.into(), exec_state.clone(), responses))?;
236    for issue in exec_state.issues() {
237        if issue.severity.is_err() {
238            return Err(ExecErrorWithState::new(
239                KclErrorWithOutputs::no_outputs(KclError::new_semantic(issue.clone().into())).into(),
240                exec_state.clone(),
241                None,
242            ));
243        }
244    }
245
246    Ok((exec_state, result.0))
247}
248
249async fn do_execute_and_snapshot(
250    ctx: &ExecutorContext,
251    program: Program,
252    deprecation_version_override: Option<&str>,
253) -> Result<(ExecState, EnvironmentRef, image::DynamicImage), ExecErrorWithState> {
254    let (exec_state, env_ref) = do_execute(ctx, program, deprecation_version_override).await?;
255    let snapshot_png_bytes = ctx
256        .prepare_snapshot()
257        .await
258        .map_err(|err| ExecErrorWithState::new(err, exec_state.clone(), None))?
259        .contents
260        .0;
261
262    // Decode the snapshot, return it.
263    let img = image::ImageReader::new(std::io::Cursor::new(snapshot_png_bytes))
264        .with_guessed_format()
265        .map_err(|e| ExecError::BadPng(e.to_string()))
266        .and_then(|x| x.decode().map_err(|e| ExecError::BadPng(e.to_string())))
267        .map_err(|err| ExecErrorWithState::new(err, exec_state.clone(), None))?;
268
269    Ok((exec_state, env_ref, img))
270}
271
272pub async fn new_context(with_auth: bool, current_file: Option<PathBuf>) -> Result<ExecutorContext, ConnectionError> {
273    new_context_with_heartbeats(with_auth, current_file, None).await
274}
275
276async fn new_context_with_heartbeats(
277    with_auth: bool,
278    current_file: Option<PathBuf>,
279    heartbeats: Option<u64>,
280) -> Result<ExecutorContext, ConnectionError> {
281    let mut client = new_zoo_client(if with_auth { None } else { Some("bad_token".to_string()) }, None)
282        .map_err(ConnectionError::CouldNotMakeClient)?;
283    if !with_auth {
284        // Use prod, don't override based on env vars.
285        // We do this so even in the engine repo, tests that need to run with
286        // no auth can fail in the same way as they would in prod.
287        client.set_base_url("https://api.zoo.dev".to_string());
288    }
289
290    let mut settings = ExecutorSettings {
291        highlight_edges: true,
292        enable_ssao: false,
293        show_grid: false,
294        replay: None,
295        project_directory: None,
296        current_file: None,
297        fixed_size_grid: true,
298        skip_artifact_graph: false,
299        heartbeats,
300        default_backface_color: Some("#00D5FF".to_owned()),
301    };
302    if let Some(current_file) = current_file {
303        settings.with_current_file(crate::TypedPath(current_file));
304    }
305    let ctx = ExecutorContext::new(&client, settings)
306        .await
307        .map_err(ConnectionError::Establishing)?;
308    Ok(ctx)
309}
310
311pub async fn execute_and_export_step(
312    code: &str,
313    current_file: Option<PathBuf>,
314) -> Result<
315    (
316        ExecState,
317        EnvironmentRef,
318        Vec<kittycad_modeling_cmds::websocket::RawFile>,
319    ),
320    ExecErrorWithState,
321> {
322    let ctx = new_context(true, current_file).await?;
323    let mut exec_state = ExecState::new(&ctx);
324    let program = Program::parse_no_errs(code).map_err(|err| {
325        ExecErrorWithState::new(KclErrorWithOutputs::no_outputs(err).into(), exec_state.clone(), None)
326    })?;
327    let result = ctx
328        .run(&program, &mut exec_state)
329        .await
330        .map_err(|err| ExecErrorWithState::new(err.into(), exec_state.clone(), None))?;
331    for issue in exec_state.issues() {
332        if issue.severity.is_err() {
333            return Err(ExecErrorWithState::new(
334                KclErrorWithOutputs::no_outputs(KclError::new_semantic(issue.clone().into())).into(),
335                exec_state.clone(),
336                None,
337            ));
338        }
339    }
340
341    let files = match ctx.export_step(true).await {
342        Ok(f) => f,
343        Err(err) => {
344            return Err(ExecErrorWithState::new(
345                ExecError::BadExport(format!("Export failed: {err:?}")),
346                exec_state.clone(),
347                None,
348            ));
349        }
350    };
351
352    ctx.close().await;
353
354    Ok((exec_state, result.0, files))
355}
356
357#[cfg(test)]
358mod tests {
359    use std::cell::Cell;
360
361    use super::*;
362
363    #[tokio::test]
364    async fn disabled_gltf_export_does_not_call_exporter() {
365        let called = Cell::new(false);
366
367        let files = export_gltf_if_requested(false, || async {
368            called.set(true);
369            Ok::<_, KclError>(Vec::new())
370        })
371        .await
372        .unwrap();
373
374        assert!(files.is_empty());
375        assert!(!called.get());
376    }
377}