datafake_rs/
generator.rs

1use crate::config::ConfigParser;
2use crate::engine::Engine;
3use crate::error::Result;
4use crate::types::{DataFakeConfig, GenerationContext};
5use serde_json::Value;
6
7pub struct DataGenerator {
8    config: DataFakeConfig,
9}
10
11impl DataGenerator {
12    pub fn new(config: DataFakeConfig) -> Self {
13        Self { config }
14    }
15
16    pub fn from_json(json_str: &str) -> Result<Self> {
17        let config = ConfigParser::parse(json_str)?;
18        Ok(Self::new(config))
19    }
20
21    pub fn from_value(json_value: Value) -> Result<Self> {
22        let config = ConfigParser::parse_value(json_value)?;
23        Ok(Self::new(config))
24    }
25
26    pub fn generate(&self) -> Result<Value> {
27        // First, convert HashMap to Map for engine
28        let variables_map: serde_json::Map<String, Value> = self
29            .config
30            .variables
31            .iter()
32            .map(|(k, v)| (k.clone(), v.clone()))
33            .collect();
34
35        // Generate all variables
36        let generated_vars = Engine::generate_variables(&variables_map)?;
37
38        // Create context with generated variables (convert back to HashMap)
39        let context = GenerationContext::with_variables(generated_vars.into_iter().collect());
40
41        // Process the schema with the context
42        Engine::process_schema(&self.config.schema, &context)
43    }
44
45    pub fn generate_batch(&self, count: usize) -> Result<Vec<Value>> {
46        let mut results = Vec::with_capacity(count);
47
48        for _ in 0..count {
49            results.push(self.generate()?);
50        }
51
52        Ok(results)
53    }
54
55    pub fn config(&self) -> &DataFakeConfig {
56        &self.config
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63    use serde_json::json;
64
65    #[test]
66    fn test_from_json() {
67        let config_json = r#"{
68            "schema": {
69                "id": {"fake": ["uuid"]},
70                "name": {"fake": ["name"]}
71            }
72        }"#;
73
74        let generator = DataGenerator::from_json(config_json).unwrap();
75        let result = generator.generate().unwrap();
76
77        assert!(result["id"].is_string());
78        assert!(result["name"].is_string());
79    }
80
81    #[test]
82    fn test_with_variables() {
83        let config_json = r#"{
84            "variables": {
85                "userId": {"fake": ["uuid"]},
86                "country": {"fake": ["country_code"]}
87            },
88            "schema": {
89                "id": {"var": "userId"},
90                "location": {
91                    "country": {"var": "country"},
92                    "city": {"fake": ["city_name"]}
93                }
94            }
95        }"#;
96
97        let generator = DataGenerator::from_json(config_json).unwrap();
98        let result = generator.generate().unwrap();
99
100        assert!(result["id"].is_string());
101        assert_eq!(result["id"], result["id"]); // Should be same for single generation
102        assert!(result["location"]["country"].is_string());
103        assert!(result["location"]["city"].is_string());
104    }
105
106    #[test]
107    fn test_generate_batch() {
108        let config_json = r#"{
109            "schema": {
110                "id": {"fake": ["uuid"]},
111                "timestamp": {"fake": ["u64"]}
112            }
113        }"#;
114
115        let generator = DataGenerator::from_json(config_json).unwrap();
116        let results = generator.generate_batch(5).unwrap();
117
118        assert_eq!(results.len(), 5);
119
120        // Each result should have unique values
121        let mut ids = std::collections::HashSet::new();
122        for result in results {
123            assert!(result["id"].is_string());
124            assert!(result["timestamp"].is_number());
125            ids.insert(result["id"].as_str().unwrap().to_string());
126        }
127        assert_eq!(ids.len(), 5); // All IDs should be unique
128    }
129
130    #[test]
131    fn test_complex_schema() {
132        let config = json!({
133            "metadata": {
134                "name": "User Profile Generator",
135                "version": "1.0.0"
136            },
137            "variables": {
138                "userId": {"fake": ["uuid"]},
139                "createdAt": {"fake": ["u64", 1000000000, 1700000000]}
140            },
141            "schema": {
142                "id": {"var": "userId"},
143                "profile": {
144                    "firstName": {"fake": ["first_name"]},
145                    "lastName": {"fake": ["last_name"]},
146                    "email": {"fake": ["email"]},
147                    "age": {"fake": ["u8", 18, 65]}
148                },
149                "address": {
150                    "street": {"fake": ["street_address"]},
151                    "city": {"fake": ["city_name"]},
152                    "zipCode": {"fake": ["zip_code"]}
153                },
154                "metadata": {
155                    "createdAt": {"var": "createdAt"},
156                    "updatedAt": {"fake": ["u64", 1700000000, 1800000000]}
157                }
158            }
159        });
160
161        let generator = DataGenerator::from_value(config).unwrap();
162        let result = generator.generate().unwrap();
163
164        // Verify structure
165        assert!(result["id"].is_string());
166        assert!(result["profile"]["firstName"].is_string());
167        assert!(result["profile"]["email"].as_str().unwrap().contains('@'));
168        assert!(result["address"]["street"].is_string());
169        assert!(result["metadata"]["createdAt"].is_number());
170
171        // Verify that variable references work
172        assert_eq!(result["id"], result["id"]);
173        assert_eq!(
174            result["metadata"]["createdAt"],
175            result["metadata"]["createdAt"]
176        );
177    }
178}