use std::collections::BTreeSet;
use crate::e2e::config::ArgMapping;
use crate::e2e::fixture::Fixture;
use super::super::helpers::is_skipped;
use super::super::test_function::{KwargRenderContext, render_kwarg_field_value};
use super::references_identifier;
pub(super) fn collect_nested_config_types(
arg: &ArgMapping,
value: &serde_json::Value,
constructor_type: Option<&str>,
context: KwargRenderContext<'_>,
used_config_types: &mut BTreeSet<String>,
) {
if let Some(obj) = value.as_object() {
for (key, field_value) in obj.iter() {
let mut nested = BTreeSet::new();
let _ = render_kwarg_field_value(
key,
field_value,
constructor_type,
&format!("/{key}"),
context,
&mut nested,
);
used_config_types.extend(nested);
}
}
if let Some(elem_type) = &arg.element_type
&& let Some(arr) = value.as_array()
{
for item in arr.iter().filter_map(|v| v.as_object()) {
for (key, field_value) in item.iter() {
let mut nested = BTreeSet::new();
let _ = render_kwarg_field_value(
key,
field_value,
Some(elem_type.as_str()),
&format!("/{key}"),
context,
&mut nested,
);
used_config_types.extend(nested);
}
}
}
}
pub(super) fn compute_pytest_and_sys_import_needs(
fixtures: &[&Fixture],
client_factory: Option<&str>,
has_error_test: bool,
is_async: bool,
) -> (bool, bool) {
let has_skipped_fixture = fixtures
.iter()
.filter(|f| !f.is_http_test())
.any(|f| is_skipped(f, "python"));
let has_pytest_skip_call = client_factory.is_some()
&& fixtures.iter().filter(|f| !f.is_http_test()).any(|f| {
let has_mock = f.mock_response.is_some() || f.http.is_some();
!has_mock && f.env.as_ref().and_then(|e| e.api_key_var.as_ref()).is_some()
});
let needs_pytest = has_error_test || is_async || has_skipped_fixture || has_pytest_skip_call;
let needs_sys_import = client_factory.is_some()
&& fixtures.iter().filter(|f| !f.is_http_test()).any(|f| {
let has_mock = f.mock_response.is_some() || f.http.is_some();
has_mock && f.env.as_ref().and_then(|e| e.api_key_var.as_ref()).is_some()
});
(needs_pytest, needs_sys_import)
}
#[derive(Clone, Copy)]
pub(super) struct ImportNeeds {
pub has_http_tests: bool,
pub needs_base64_import: bool,
pub needs_json_import: bool,
pub needs_os_import: bool,
pub needs_path_import: bool,
pub needs_sys_import: bool,
pub needs_pytest: bool,
}
pub(super) fn finalize_stdlib_and_bare_imports(
fixtures_body: &str,
needs: ImportNeeds,
stdlib_imports: &mut Vec<String>,
thirdparty_bare: &mut Vec<String>,
) {
let needs_json_import = needs.needs_json_import
|| references_identifier(fixtures_body, "json.dumps")
|| references_identifier(fixtures_body, "json.loads");
let needs_re_import =
references_identifier(fixtures_body, "re.match") || references_identifier(fixtures_body, "re.search");
if needs.needs_base64_import {
stdlib_imports.push("import base64".to_string());
}
if needs_json_import {
stdlib_imports.push("import json".to_string());
}
if needs.needs_os_import {
stdlib_imports.push("import os".to_string());
}
if needs.needs_path_import {
stdlib_imports.push("from pathlib import Path".to_string());
}
if needs_re_import {
stdlib_imports.push("import re".to_string());
}
if needs.has_http_tests {
stdlib_imports.push("import urllib.request".to_string());
}
if needs.needs_sys_import {
stdlib_imports.push("import sys".to_string());
}
if needs.needs_pytest {
thirdparty_bare.push("import pytest".to_string());
}
stdlib_imports.sort_by(|a, b| (a.starts_with("from "), a).cmp(&(b.starts_with("from "), b)));
thirdparty_bare.sort();
}
pub(super) fn prune_unreferenced_from_imports(imports: &mut Vec<String>, emitted: &[&str]) {
let pruned: Vec<String> = imports
.iter()
.filter_map(|line| {
let Some((prefix, names)) = line.split_once(" import ") else {
return Some(line.clone());
};
let kept: Vec<&str> = names
.split(", ")
.map(str::trim)
.filter(|name| emitted.iter().any(|source| references_identifier(source, name)))
.collect();
if kept.is_empty() {
return None;
}
Some(format!("{prefix} import {}", kept.join(", ")))
})
.collect();
*imports = pruned;
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::ir::{FieldDef, TypeDef, TypeRef};
use crate::e2e::codegen::python::test_function::LeafSource;
use crate::e2e::fixture::{Assertion, Fixture};
use std::collections::HashMap;
fn config_arg() -> ArgMapping {
ArgMapping {
name: "config".to_string(),
field: "input.config".to_string(),
arg_type: "json_object".to_string(),
optional: false,
owned: false,
element_type: None,
go_type: None,
vec_inner_is_ref: false,
trait_name: None,
}
}
#[test]
fn collect_nested_config_types_collects_single_object_nested_type() {
let outer = TypeDef {
name: "ExtractionConfig".to_string(),
rust_path: "demo::ExtractionConfig".to_string(),
fields: vec![FieldDef {
name: "nested".to_string(),
ty: TypeRef::Named("NestedConfig".to_string()),
..Default::default()
}],
..Default::default()
};
let inner = TypeDef {
name: "NestedConfig".to_string(),
rust_path: "demo::NestedConfig".to_string(),
fields: vec![FieldDef {
name: "value".to_string(),
ty: TypeRef::String,
..Default::default()
}],
..Default::default()
};
let type_defs = vec![outer, inner];
let arg = config_arg();
let value = serde_json::json!({"nested": {"value": "x"}});
let mut used_config_types: BTreeSet<String> = BTreeSet::new();
let context = KwargRenderContext {
type_defs: &type_defs,
enums: &[],
enum_fields: &HashMap::new(),
docs_files: &[],
leaf_source: LeafSource::Literal,
};
collect_nested_config_types(&arg, &value, Some("ExtractionConfig"), context, &mut used_config_types);
assert_eq!(
used_config_types,
["NestedConfig".to_string()].into_iter().collect(),
"the nested struct type must be collected for import, got: {used_config_types:?}"
);
}
#[test]
fn collect_nested_config_types_collects_map_value_nested_type() {
let outer = TypeDef {
name: "ExtractionConfig".to_string(),
rust_path: "demo::ExtractionConfig".to_string(),
fields: vec![FieldDef {
name: "profiles".to_string(),
ty: TypeRef::Map(
Box::new(TypeRef::String),
Box::new(TypeRef::Named("NestedConfig".to_string())),
),
..Default::default()
}],
..Default::default()
};
let inner = TypeDef {
name: "NestedConfig".to_string(),
rust_path: "demo::NestedConfig".to_string(),
fields: vec![FieldDef {
name: "value".to_string(),
ty: TypeRef::String,
..Default::default()
}],
..Default::default()
};
let type_defs = vec![outer, inner];
let arg = config_arg();
let value = serde_json::json!({"profiles": {"first": {"value": "x"}}});
let mut used_config_types: BTreeSet<String> = BTreeSet::new();
let context = KwargRenderContext {
type_defs: &type_defs,
enums: &[],
enum_fields: &HashMap::new(),
docs_files: &[],
leaf_source: LeafSource::Literal,
};
collect_nested_config_types(&arg, &value, Some("ExtractionConfig"), context, &mut used_config_types);
assert_eq!(
used_config_types,
["NestedConfig".to_string()].into_iter().collect(),
"the map value's struct type must be collected for import, got: {used_config_types:?}"
);
}
#[test]
fn collect_nested_config_types_collects_batch_element_nested_type() {
let item_type = TypeDef {
name: "BatchFileItem".to_string(),
rust_path: "demo::BatchFileItem".to_string(),
fields: vec![FieldDef {
name: "nested".to_string(),
ty: TypeRef::Named("NestedConfig".to_string()),
..Default::default()
}],
..Default::default()
};
let inner = TypeDef {
name: "NestedConfig".to_string(),
rust_path: "demo::NestedConfig".to_string(),
fields: vec![FieldDef {
name: "value".to_string(),
ty: TypeRef::String,
..Default::default()
}],
..Default::default()
};
let type_defs = vec![item_type, inner];
let mut arg = config_arg();
arg.element_type = Some("BatchFileItem".to_string());
let value = serde_json::json!([{"nested": {"value": "x"}}]);
let mut used_config_types: BTreeSet<String> = BTreeSet::new();
let context = KwargRenderContext {
type_defs: &type_defs,
enums: &[],
enum_fields: &HashMap::new(),
docs_files: &[],
leaf_source: LeafSource::Literal,
};
collect_nested_config_types(&arg, &value, None, context, &mut used_config_types);
assert_eq!(
used_config_types,
["NestedConfig".to_string()].into_iter().collect(),
"the batch element's nested struct type must be collected for import, got: {used_config_types:?}"
);
}
#[test]
fn prune_unreferenced_from_imports_drops_a_collected_but_unreferenced_type() {
let outer = TypeDef {
name: "ExtractionConfig".to_string(),
rust_path: "demo::ExtractionConfig".to_string(),
fields: vec![FieldDef {
name: "nested".to_string(),
ty: TypeRef::Named("GhostConfig".to_string()),
..Default::default()
}],
..Default::default()
};
let inner = TypeDef {
name: "GhostConfig".to_string(),
rust_path: "demo::GhostConfig".to_string(),
fields: vec![FieldDef {
name: "value".to_string(),
ty: TypeRef::String,
..Default::default()
}],
..Default::default()
};
let type_defs = vec![outer, inner];
let arg = config_arg();
let value = serde_json::json!({"nested": {"value": "x"}});
let mut used_config_types: BTreeSet<String> = BTreeSet::new();
let context = KwargRenderContext {
type_defs: &type_defs,
enums: &[],
enum_fields: &HashMap::new(),
docs_files: &[],
leaf_source: LeafSource::Literal,
};
collect_nested_config_types(&arg, &value, Some("ExtractionConfig"), context, &mut used_config_types);
assert!(
used_config_types.contains("GhostConfig"),
"test setup: GhostConfig must be collected as a candidate for this to be a real test of pruning"
);
let mut imports = vec!["from sample_pkg import process, ExtractionConfig, GhostConfig".to_string()];
let emitted_body = " result = process(opts=ExtractionConfig(nested={\"value\": \"x\"}))\n";
prune_unreferenced_from_imports(&mut imports, &[emitted_body]);
assert_eq!(
imports,
vec!["from sample_pkg import process, ExtractionConfig".to_string()],
"a collected-but-unreferenced type must not survive into the import line, got: {imports:?}"
);
}
fn fixture_with_input(id: &str, input: serde_json::Value) -> Fixture {
Fixture {
docs: None,
requirements: Vec::new(),
id: id.to_string(),
description: "Smoke test".to_string(),
input,
http: None,
asyncapi: None,
websocket: None,
preserve_input_urls: false,
assertions: vec![Assertion {
assertion_type: "not_error".to_string(),
..Default::default()
}],
call: None,
skip: None,
env: None,
setup: Vec::new(),
visitor: None,
args: vec![],
assertion_recipes: vec![],
mock_response: None,
source: String::new(),
category: None,
tags: Vec::new(),
}
}
fn e2e_config_with_options_type(options_type: &str) -> crate::e2e::config::E2eConfig {
let mut e2e_config = crate::e2e::config::E2eConfig::default();
e2e_config.call.module = "sample_pkg".to_string();
e2e_config.call.function = "process".to_string();
e2e_config.call.args = vec![config_arg()];
e2e_config.call.overrides.insert(
"python".to_string(),
crate::e2e::config::CallOverride {
options_type: Some(options_type.to_string()),
..Default::default()
},
);
e2e_config
}
#[test]
fn render_test_file_imports_nested_config_class_for_object_field() {
let outer = TypeDef {
name: "ExtractionConfig".to_string(),
rust_path: "demo::ExtractionConfig".to_string(),
fields: vec![FieldDef {
name: "nested".to_string(),
ty: TypeRef::Named("NestedConfig".to_string()),
..Default::default()
}],
..Default::default()
};
let inner = TypeDef {
name: "NestedConfig".to_string(),
rust_path: "demo::NestedConfig".to_string(),
fields: vec![FieldDef {
name: "value".to_string(),
ty: TypeRef::String,
..Default::default()
}],
..Default::default()
};
let type_defs = vec![outer, inner];
let e2e_config = e2e_config_with_options_type("ExtractionConfig");
let config = crate::core::config::ResolvedCrateConfig::default();
let fixture = fixture_with_input(
"nested_object",
serde_json::json!({"config": {"nested": {"value": "x"}}}),
);
let fixtures: Vec<&Fixture> = vec![&fixture];
let out = super::super::render_test_file(
"smoke",
&fixtures,
&e2e_config,
&config,
&type_defs,
&[],
&[],
&[],
false,
);
let import_line = out
.lines()
.find(|line| line.starts_with("from sample_pkg import"))
.unwrap_or_else(|| panic!("no `from sample_pkg import ...` line in output, got:\n{out}"));
assert!(
import_line.contains("NestedConfig"),
"the nested config class must be imported, got: {import_line:?}"
);
assert!(
out.contains("NestedConfig(value="),
"the nested config class must be constructed, got:\n{out}"
);
}
#[test]
fn render_test_file_imports_nested_config_class_at_depth() {
let outer = TypeDef {
name: "ExtractionConfig".to_string(),
rust_path: "demo::ExtractionConfig".to_string(),
fields: vec![FieldDef {
name: "nested".to_string(),
ty: TypeRef::Named("NestedConfig".to_string()),
..Default::default()
}],
..Default::default()
};
let middle = TypeDef {
name: "NestedConfig".to_string(),
rust_path: "demo::NestedConfig".to_string(),
fields: vec![FieldDef {
name: "inner".to_string(),
ty: TypeRef::Named("DeeperConfig".to_string()),
..Default::default()
}],
..Default::default()
};
let deepest = TypeDef {
name: "DeeperConfig".to_string(),
rust_path: "demo::DeeperConfig".to_string(),
fields: vec![FieldDef {
name: "value".to_string(),
ty: TypeRef::String,
..Default::default()
}],
..Default::default()
};
let type_defs = vec![outer, middle, deepest];
let e2e_config = e2e_config_with_options_type("ExtractionConfig");
let config = crate::core::config::ResolvedCrateConfig::default();
let fixture = fixture_with_input(
"nested_depth",
serde_json::json!({"config": {"nested": {"inner": {"value": "x"}}}}),
);
let fixtures: Vec<&Fixture> = vec![&fixture];
let out = super::super::render_test_file(
"smoke",
&fixtures,
&e2e_config,
&config,
&type_defs,
&[],
&[],
&[],
false,
);
let import_line = out
.lines()
.find(|line| line.starts_with("from sample_pkg import"))
.unwrap_or_else(|| panic!("no `from sample_pkg import ...` line in output, got:\n{out}"));
assert!(
import_line.contains("NestedConfig") && import_line.contains("DeeperConfig"),
"both nested levels must be imported, got: {import_line:?}"
);
assert!(
out.contains("NestedConfig(inner=DeeperConfig(value="),
"both levels must be constructed, got:\n{out}"
);
}
#[test]
fn render_test_file_imports_nested_config_class_for_map_field() {
let outer = TypeDef {
name: "ExtractionConfig".to_string(),
rust_path: "demo::ExtractionConfig".to_string(),
fields: vec![FieldDef {
name: "profiles".to_string(),
ty: TypeRef::Map(
Box::new(TypeRef::String),
Box::new(TypeRef::Named("NestedConfig".to_string())),
),
..Default::default()
}],
..Default::default()
};
let inner = TypeDef {
name: "NestedConfig".to_string(),
rust_path: "demo::NestedConfig".to_string(),
fields: vec![FieldDef {
name: "value".to_string(),
ty: TypeRef::String,
..Default::default()
}],
..Default::default()
};
let type_defs = vec![outer, inner];
let e2e_config = e2e_config_with_options_type("ExtractionConfig");
let config = crate::core::config::ResolvedCrateConfig::default();
let fixture = fixture_with_input(
"nested_map",
serde_json::json!({"config": {"profiles": {"first": {"value": "x"}}}}),
);
let fixtures: Vec<&Fixture> = vec![&fixture];
let out = super::super::render_test_file(
"smoke",
&fixtures,
&e2e_config,
&config,
&type_defs,
&[],
&[],
&[],
false,
);
let import_line = out
.lines()
.find(|line| line.starts_with("from sample_pkg import"))
.unwrap_or_else(|| panic!("no `from sample_pkg import ...` line in output, got:\n{out}"));
assert!(
import_line.contains("NestedConfig"),
"the map value's class must be imported, got: {import_line:?}"
);
assert!(
out.contains(r#"{"first": NestedConfig(value="#),
"the map value must be constructed, got:\n{out}"
);
}
#[test]
fn render_test_file_constructs_nested_config_class_for_mock_url_fixture() {
let outer = TypeDef {
name: "ExtractionConfig".to_string(),
rust_path: "demo::ExtractionConfig".to_string(),
fields: vec![FieldDef {
name: "nested".to_string(),
ty: TypeRef::Named("NestedConfig".to_string()),
..Default::default()
}],
..Default::default()
};
let inner = TypeDef {
name: "NestedConfig".to_string(),
rust_path: "demo::NestedConfig".to_string(),
fields: vec![FieldDef {
name: "url".to_string(),
ty: TypeRef::String,
..Default::default()
}],
..Default::default()
};
let type_defs = vec![outer, inner];
let e2e_config = e2e_config_with_options_type("ExtractionConfig");
let config = crate::core::config::ResolvedCrateConfig::default();
let fixture = fixture_with_input(
"nested_mock_url",
serde_json::json!({"config": {"nested": {"url": "$mock_url/path"}}}),
);
let fixtures: Vec<&Fixture> = vec![&fixture];
let out = super::super::render_test_file(
"smoke",
&fixtures,
&e2e_config,
&config,
&type_defs,
&[],
&[],
&[],
false,
);
assert!(
out.contains("NestedConfig(url="),
"the nested struct field must still be constructed with its own class for a \
$mock_url fixture, got:\n{out}"
);
assert!(
!out.contains("**json.loads"),
"the nested constructor must not fall back to unpacking a raw dict, got:\n{out}"
);
}
}