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 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 let generated_vars = Engine::generate_variables(&variables_map)?;
37
38 let context = GenerationContext::with_variables(generated_vars.into_iter().collect());
40
41 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"]); 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 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); }
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 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 assert_eq!(result["id"], result["id"]);
173 assert_eq!(
174 result["metadata"]["createdAt"],
175 result["metadata"]["createdAt"]
176 );
177 }
178}