Skip to main content

previa_engine/
extractions.rs

1use std::collections::{HashMap, HashSet};
2
3use regex::Regex;
4use serde_json::Value;
5
6use crate::core::types::{PipelineStep, StepExecutionResult};
7
8pub fn validate_step_extractions(step: &PipelineStep) -> Vec<String> {
9    let mut names = HashSet::new();
10    let mut errors = Vec::new();
11
12    for extraction in &step.extracts {
13        if !is_valid_name(&extraction.name) {
14            errors.push(format!(
15                "extraction '{}' has an invalid name; use lowercase letters, digits, '-' or '_'",
16                extraction.name
17            ));
18        }
19        if !names.insert(extraction.name.as_str()) {
20            errors.push(format!(
21                "extraction '{}' is a duplicate within step '{}'",
22                extraction.name, step.id
23            ));
24        }
25        if extraction.field != "body"
26            && (!extraction.field.starts_with("body.") || extraction.field == "body.")
27        {
28            errors.push(format!(
29                "extraction '{}' field must be 'body' or start with 'body.'",
30                extraction.name
31            ));
32        }
33
34        match Regex::new(&extraction.regex) {
35            Ok(regex) if extraction.group >= regex.captures_len() => errors.push(format!(
36                "extraction '{}' group {} does not exist in regex",
37                extraction.name, extraction.group
38            )),
39            Ok(_) => {}
40            Err(_) => errors.push(format!(
41                "extraction '{}' has an invalid regex",
42                extraction.name
43            )),
44        }
45    }
46
47    errors
48}
49
50pub fn evaluate_step_extractions(
51    step: &PipelineStep,
52    result: &StepExecutionResult,
53) -> Result<HashMap<String, String>, String> {
54    let validation_errors = validate_step_extractions(step);
55    if !validation_errors.is_empty() {
56        return Err(validation_errors.join("; "));
57    }
58
59    let mut values = HashMap::new();
60    for extraction in &step.extracts {
61        let captured = resolve_source(&extraction.field, result).and_then(|source| {
62            let regex = Regex::new(&extraction.regex).ok()?;
63            regex
64                .captures(&source)
65                .and_then(|captures| captures.get(extraction.group))
66                .map(|capture| capture.as_str().to_owned())
67        });
68
69        match captured {
70            Some(value) => {
71                values.insert(extraction.name.clone(), value);
72            }
73            None if extraction.required => {
74                return Err(format!(
75                    "required extraction '{}' did not produce a value",
76                    extraction.name
77                ));
78            }
79            None => {}
80        }
81    }
82
83    Ok(values)
84}
85
86fn is_valid_name(name: &str) -> bool {
87    !name.is_empty()
88        && name
89            .chars()
90            .all(|ch| ch.is_ascii_lowercase() || ch.is_ascii_digit() || ch == '-' || ch == '_')
91}
92
93fn resolve_source(field: &str, result: &StepExecutionResult) -> Option<String> {
94    let response = result.response.as_ref()?;
95    let value = if field == "body" {
96        &response.body
97    } else {
98        resolve_json_path(&response.body, field.strip_prefix("body.")?)?
99    };
100    value_to_string(value)
101}
102
103fn resolve_json_path<'a>(value: &'a Value, path: &str) -> Option<&'a Value> {
104    let mut current = value;
105    for segment in path.split('.') {
106        current = match current {
107            Value::Object(map) => map.get(segment)?,
108            Value::Array(items) => items.get(segment.parse::<usize>().ok()?)?,
109            _ => return None,
110        };
111    }
112    Some(current)
113}
114
115fn value_to_string(value: &Value) -> Option<String> {
116    match value {
117        Value::String(value) => Some(value.clone()),
118        Value::Number(value) => Some(value.to_string()),
119        Value::Bool(value) => Some(value.to_string()),
120        Value::Null | Value::Array(_) | Value::Object(_) => None,
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use std::collections::HashMap;
127
128    use serde_json::json;
129
130    use crate::{
131        PipelineStep, StepExecutionResult, StepExtraction, StepResponse, evaluate_step_extractions,
132        validate_step_extractions,
133    };
134
135    fn step(extracts: Vec<StepExtraction>) -> PipelineStep {
136        PipelineStep {
137            id: "email".to_owned(),
138            name: "Read e-mail".to_owned(),
139            description: None,
140            method: "GET".to_owned(),
141            url: "https://example.test/message".to_owned(),
142            headers: HashMap::new(),
143            body: None,
144            operation_id: None,
145            delay: None,
146            retry: None,
147            asserts: Vec::new(),
148            extracts,
149        }
150    }
151
152    fn result(body: serde_json::Value) -> StepExecutionResult {
153        StepExecutionResult {
154            step_id: "email".to_owned(),
155            status: "success".to_owned(),
156            request: None,
157            response: Some(StepResponse {
158                status: 200,
159                status_text: "OK".to_owned(),
160                headers: HashMap::new(),
161                body,
162            }),
163            error: None,
164            duration: Some(1),
165            attempts: None,
166            attempt: Some(1),
167            max_attempts: Some(1),
168            assert_results: None,
169            extracts: HashMap::new(),
170        }
171    }
172
173    fn extraction(name: &str, field: &str, regex: &str) -> StepExtraction {
174        StepExtraction {
175            name: name.to_owned(),
176            field: field.to_owned(),
177            regex: regex.to_owned(),
178            group: 1,
179            required: true,
180        }
181    }
182
183    #[test]
184    fn extracts_capture_from_nested_json_string() {
185        let step = step(vec![extraction(
186            "code",
187            "body.HTML",
188            r"<strong>[[:space:]]*([0-9]{6})[[:space:]]*</strong>",
189        )]);
190        let result = result(json!({"HTML": "<p><strong>123456</strong></p>"}));
191
192        assert_eq!(
193            evaluate_step_extractions(&step, &result)
194                .expect("capture should succeed")
195                .get("code"),
196            Some(&"123456".to_owned())
197        );
198    }
199
200    #[test]
201    fn group_zero_extracts_the_entire_match_from_string_body() {
202        let mut definition = extraction("code", "body", r"[0-9]{6}");
203        definition.group = 0;
204        let step = step(vec![definition]);
205        let result = result(json!("Login code: 123456"));
206
207        assert_eq!(
208            evaluate_step_extractions(&step, &result)
209                .expect("capture should succeed")
210                .get("code"),
211            Some(&"123456".to_owned())
212        );
213    }
214
215    #[test]
216    fn missing_optional_capture_is_omitted() {
217        let mut definition = extraction("code", "body.HTML", r"([0-9]{6})");
218        definition.required = false;
219        let step = step(vec![definition]);
220        let result = result(json!({"HTML": "no code"}));
221
222        assert!(
223            evaluate_step_extractions(&step, &result)
224                .expect("optional capture should not fail")
225                .is_empty()
226        );
227    }
228
229    #[test]
230    fn missing_required_capture_fails_without_response_content() {
231        let step = step(vec![extraction("code", "body.HTML", r"([0-9]{6})")]);
232        let result = result(json!({"HTML": "sensitive message without code"}));
233
234        let error =
235            evaluate_step_extractions(&step, &result).expect_err("required capture should fail");
236
237        assert!(error.contains("code"));
238        assert!(!error.contains("sensitive message"));
239    }
240
241    #[test]
242    fn validates_invalid_regex_duplicate_and_invalid_names() {
243        let step = step(vec![
244            extraction("bad name", "body.HTML", "("),
245            extraction("bad name", "body.HTML", r"([0-9]{6})"),
246        ]);
247
248        let errors = validate_step_extractions(&step);
249
250        assert!(errors.iter().any(|error| error.contains("invalid name")));
251        assert!(errors.iter().any(|error| error.contains("invalid regex")));
252        assert!(errors.iter().any(|error| error.contains("duplicate")));
253    }
254
255    #[test]
256    fn validates_source_path_and_capture_group() {
257        let invalid_path = extraction("code", "header.subject", r"([0-9]{6})");
258        let mut invalid_group = extraction("token", "body.Text", r"([a-z]+)");
259        invalid_group.group = 2;
260        let step = step(vec![invalid_path, invalid_group]);
261
262        let errors = validate_step_extractions(&step);
263
264        assert!(errors.iter().any(|error| error.contains("field")));
265        assert!(errors.iter().any(|error| error.contains("group")));
266    }
267}