Skip to main content

_standard_knowledge_py/
test_suite.rs

1use pyo3::prelude::*;
2use pyo3::types::PyDict;
3use standard_knowledge::qartod::{types::ArgumentValue, TestSuite};
4use std::collections::HashMap;
5
6#[pyclass(name = "TestSuite")]
7#[derive(Clone)]
8pub struct PyTestSuite {
9    test_suite: Box<dyn TestSuite>,
10}
11
12impl PyTestSuite {
13    pub fn new(test_suite: Box<dyn TestSuite>) -> Self {
14        Self { test_suite }
15    }
16}
17
18/// A QARTOD test suite
19#[pymethods]
20impl PyTestSuite {
21    fn __repr__(&self) -> PyResult<String> {
22        let info = self.test_suite.info();
23        Ok(format!("<TestSuite: {} ({})>", info.name, info.slug))
24    }
25
26    /// Get test suite information
27    fn info(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
28        let info = self.test_suite.info();
29        let dict = PyDict::new(py);
30        dict.set_item("name", info.name)?;
31        dict.set_item("slug", info.slug)?;
32        dict.set_item("summary", info.summary)?;
33        dict.set_item("description", info.description)?;
34        // TODO: Add arguments and test_types if needed
35        Ok(dict.into())
36    }
37
38    /// Generate a configuration for the test suite
39    fn scaffold(
40        &self,
41        py: Python<'_>,
42        arguments: HashMap<String, Py<PyAny>>,
43    ) -> PyResult<Py<PyAny>> {
44        // Convert Python arguments to ArgumentValue
45        let mut rust_args = HashMap::new();
46        for (key, value) in arguments {
47            // For now, we'll handle basic types. This could be expanded.
48            if let Ok(s) = value.extract::<String>(py) {
49                rust_args.insert(key, ArgumentValue::String(s));
50            } else if let Ok(f) = value.extract::<f64>(py) {
51                rust_args.insert(key, ArgumentValue::Float(f));
52            } else if let Ok(i) = value.extract::<i64>(py) {
53                rust_args.insert(key, ArgumentValue::Int(i));
54            }
55            // Add more types as needed
56        }
57
58        match self.test_suite.scaffold(rust_args) {
59            Ok(config) => {
60                // Convert ConfigStream to Python dict
61                let result = PyDict::new(py);
62
63                // Convert the qartod config to a Python dict
64                let qartod_config = &config.qartod;
65                let qartod_dict = PyDict::new(py);
66
67                if let Some(ref gross_range) = qartod_config.gross_range_test {
68                    let test_dict = PyDict::new(py);
69                    test_dict.set_item(
70                        "fail_span",
71                        vec![gross_range.fail_span.0, gross_range.fail_span.1],
72                    )?;
73                    test_dict.set_item(
74                        "suspect_span",
75                        vec![gross_range.suspect_span.0, gross_range.suspect_span.1],
76                    )?;
77                    qartod_dict.set_item("gross_range_test", test_dict)?;
78                }
79
80                if let Some(ref flat_line) = qartod_config.flat_line_test {
81                    let test_dict = PyDict::new(py);
82                    test_dict.set_item("tolerance", flat_line.tolerance)?;
83                    test_dict.set_item("suspect_threshold", flat_line.suspect_threshold)?;
84                    test_dict.set_item("fail_threshold", flat_line.fail_threshold)?;
85                    qartod_dict.set_item("flat_line_test", test_dict)?;
86                }
87
88                if let Some(ref spike) = qartod_config.spike_test {
89                    let test_dict = PyDict::new(py);
90                    test_dict.set_item("suspect_threshold", spike.suspect_threshold)?;
91                    test_dict.set_item("fail_threshold", spike.fail_threshold)?;
92                    qartod_dict.set_item("spike_test", test_dict)?;
93                }
94
95                if let Some(ref rate_of_change) = qartod_config.rate_of_change_test {
96                    let test_dict = PyDict::new(py);
97                    test_dict.set_item("threshold", rate_of_change.threshold)?;
98                    qartod_dict.set_item("rate_of_change_test", test_dict)?;
99                }
100
101                result.set_item("qartod", qartod_dict)?;
102
103                Ok(result.into())
104            }
105            Err(e) => Err(PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e)),
106        }
107    }
108}