datasynth_graph/exporters/
rustgraph.rs

1//! RustGraph JSON exporter.
2//!
3//! Exports graph data in JSON format compatible with RustGraph/RustAssureTwin:
4//! - JSONL files for nodes and edges (streaming-friendly)
5//! - JSON array format for batch import
6//! - Full temporal and feature metadata
7
8use std::collections::HashMap;
9use std::fs::{self, File};
10use std::io::{BufWriter, Write};
11use std::path::Path;
12
13use chrono::{DateTime, NaiveDateTime, Utc};
14use serde::{Deserialize, Serialize};
15use serde_json::Value;
16
17use crate::models::{EdgeDirection, EdgeProperty, Graph, GraphEdge, GraphNode, NodeProperty};
18
19/// Configuration for RustGraph export.
20#[derive(Debug, Clone)]
21pub struct RustGraphExportConfig {
22    /// Include numeric features in output.
23    pub include_features: bool,
24    /// Include temporal metadata (valid_from, valid_to, transaction_time).
25    pub include_temporal: bool,
26    /// Include ML labels in output.
27    pub include_labels: bool,
28    /// Source name for provenance tracking.
29    pub source_name: String,
30    /// Optional batch ID for grouping exports.
31    pub batch_id: Option<String>,
32    /// Output format (JsonLines or JsonArray).
33    pub output_format: RustGraphOutputFormat,
34    /// Export node properties.
35    pub export_node_properties: bool,
36    /// Export edge properties.
37    pub export_edge_properties: bool,
38    /// Pretty print JSON (for debugging).
39    pub pretty_print: bool,
40}
41
42impl Default for RustGraphExportConfig {
43    fn default() -> Self {
44        Self {
45            include_features: true,
46            include_temporal: true,
47            include_labels: true,
48            source_name: "datasynth".to_string(),
49            batch_id: None,
50            output_format: RustGraphOutputFormat::JsonLines,
51            export_node_properties: true,
52            export_edge_properties: true,
53            pretty_print: false,
54        }
55    }
56}
57
58/// Output format for RustGraph export.
59#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
60#[serde(rename_all = "snake_case")]
61pub enum RustGraphOutputFormat {
62    /// JSON Lines format (one JSON object per line).
63    JsonLines,
64    /// JSON array format (single array containing all objects).
65    JsonArray,
66}
67
68/// Node output compatible with RustGraph CreateNodeRequest.
69#[derive(Debug, Clone, Serialize, Deserialize)]
70pub struct RustGraphNodeOutput {
71    /// Node type (e.g., "Account", "Vendor", "Customer").
72    pub node_type: String,
73    /// Unique identifier.
74    pub id: String,
75    /// Node properties as key-value pairs.
76    pub properties: HashMap<String, Value>,
77    /// Metadata for the node.
78    pub metadata: RustGraphNodeMetadata,
79}
80
81/// Metadata for a RustGraph node.
82#[derive(Debug, Clone, Serialize, Deserialize)]
83pub struct RustGraphNodeMetadata {
84    /// Source system identifier.
85    pub source: String,
86    /// Generation timestamp.
87    pub generated_at: DateTime<Utc>,
88    /// Valid time start (business time).
89    #[serde(skip_serializing_if = "Option::is_none")]
90    pub valid_from: Option<NaiveDateTime>,
91    /// Valid time end (business time).
92    #[serde(skip_serializing_if = "Option::is_none")]
93    pub valid_to: Option<NaiveDateTime>,
94    /// Transaction time (system time).
95    #[serde(skip_serializing_if = "Option::is_none")]
96    pub transaction_time: Option<DateTime<Utc>>,
97    /// Custom labels for classification.
98    pub labels: HashMap<String, String>,
99    /// Numeric feature vector for ML.
100    #[serde(skip_serializing_if = "Option::is_none")]
101    pub features: Option<Vec<f64>>,
102    /// Batch identifier for grouping.
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub batch_id: Option<String>,
105    /// Whether this node is an anomaly.
106    #[serde(default)]
107    pub is_anomaly: bool,
108    /// Anomaly type if anomalous.
109    #[serde(skip_serializing_if = "Option::is_none")]
110    pub anomaly_type: Option<String>,
111}
112
113/// Edge output compatible with RustGraph CreateEdgeRequest.
114#[derive(Debug, Clone, Serialize, Deserialize)]
115pub struct RustGraphEdgeOutput {
116    /// Edge type (e.g., "Transaction", "Approval", "Ownership").
117    pub edge_type: String,
118    /// Unique identifier.
119    pub id: String,
120    /// Source node ID.
121    pub source_id: String,
122    /// Target node ID.
123    pub target_id: String,
124    /// Edge properties as key-value pairs.
125    pub properties: HashMap<String, Value>,
126    /// Metadata for the edge.
127    pub metadata: RustGraphEdgeMetadata,
128}
129
130/// Metadata for a RustGraph edge.
131#[derive(Debug, Clone, Serialize, Deserialize)]
132pub struct RustGraphEdgeMetadata {
133    /// Source system identifier.
134    pub source: String,
135    /// Generation timestamp.
136    pub generated_at: DateTime<Utc>,
137    /// Edge weight (e.g., transaction amount).
138    #[serde(skip_serializing_if = "Option::is_none")]
139    pub weight: Option<f64>,
140    /// Valid time start (business time).
141    #[serde(skip_serializing_if = "Option::is_none")]
142    pub valid_from: Option<NaiveDateTime>,
143    /// Valid time end (business time).
144    #[serde(skip_serializing_if = "Option::is_none")]
145    pub valid_to: Option<NaiveDateTime>,
146    /// Custom labels for classification.
147    pub labels: HashMap<String, String>,
148    /// Numeric feature vector for ML.
149    #[serde(skip_serializing_if = "Option::is_none")]
150    pub features: Option<Vec<f64>>,
151    /// Whether edge is directed.
152    pub is_directed: bool,
153    /// Whether this edge is an anomaly.
154    #[serde(default)]
155    pub is_anomaly: bool,
156    /// Anomaly type if anomalous.
157    #[serde(skip_serializing_if = "Option::is_none")]
158    pub anomaly_type: Option<String>,
159    /// Batch identifier for grouping.
160    #[serde(skip_serializing_if = "Option::is_none")]
161    pub batch_id: Option<String>,
162}
163
164/// Metadata about the exported RustGraph data.
165#[derive(Debug, Clone, Serialize, Deserialize)]
166pub struct RustGraphMetadata {
167    /// Graph name.
168    pub name: String,
169    /// Number of nodes exported.
170    pub num_nodes: usize,
171    /// Number of edges exported.
172    pub num_edges: usize,
173    /// Node type counts.
174    pub node_type_counts: HashMap<String, usize>,
175    /// Edge type counts.
176    pub edge_type_counts: HashMap<String, usize>,
177    /// Node feature dimension (0 if no features).
178    pub node_feature_dim: usize,
179    /// Edge feature dimension (0 if no features).
180    pub edge_feature_dim: usize,
181    /// Graph density (edges / possible edges).
182    pub graph_density: f64,
183    /// Number of anomalous nodes.
184    pub anomalous_nodes: usize,
185    /// Number of anomalous edges.
186    pub anomalous_edges: usize,
187    /// Source system identifier.
188    pub source: String,
189    /// Generation timestamp.
190    pub generated_at: DateTime<Utc>,
191    /// Output format used.
192    pub output_format: String,
193    /// Files included in export.
194    pub files: Vec<String>,
195}
196
197/// RustGraph JSON exporter.
198pub struct RustGraphExporter {
199    config: RustGraphExportConfig,
200}
201
202impl RustGraphExporter {
203    /// Creates a new RustGraph exporter with the given configuration.
204    pub fn new(config: RustGraphExportConfig) -> Self {
205        Self { config }
206    }
207
208    /// Exports a graph to RustGraph JSON format.
209    pub fn export(&self, graph: &Graph, output_dir: &Path) -> std::io::Result<RustGraphMetadata> {
210        fs::create_dir_all(output_dir)?;
211
212        let mut files = Vec::new();
213        let generated_at = Utc::now();
214
215        // Export nodes
216        let (node_type_counts, anomalous_nodes, node_feature_dim) =
217            self.export_nodes(graph, output_dir, &mut files, generated_at)?;
218
219        // Export edges
220        let (edge_type_counts, anomalous_edges, edge_feature_dim) =
221            self.export_edges(graph, output_dir, &mut files, generated_at)?;
222
223        // Calculate graph density
224        let n = graph.node_count();
225        let possible_edges = if n > 1 { n * (n - 1) } else { 1 };
226        let graph_density = graph.edge_count() as f64 / possible_edges as f64;
227
228        // Create metadata
229        let metadata = RustGraphMetadata {
230            name: graph.name.clone(),
231            num_nodes: graph.node_count(),
232            num_edges: graph.edge_count(),
233            node_type_counts,
234            edge_type_counts,
235            node_feature_dim,
236            edge_feature_dim,
237            graph_density,
238            anomalous_nodes,
239            anomalous_edges,
240            source: self.config.source_name.clone(),
241            generated_at,
242            output_format: match self.config.output_format {
243                RustGraphOutputFormat::JsonLines => "jsonl".to_string(),
244                RustGraphOutputFormat::JsonArray => "json".to_string(),
245            },
246            files: files.clone(),
247        };
248
249        // Write metadata
250        let metadata_path = output_dir.join("metadata.json");
251        let file = File::create(&metadata_path)?;
252        if self.config.pretty_print {
253            serde_json::to_writer_pretty(file, &metadata)?;
254        } else {
255            serde_json::to_writer(file, &metadata)?;
256        }
257        files.push("metadata.json".to_string());
258
259        Ok(metadata)
260    }
261
262    /// Exports nodes to file(s).
263    fn export_nodes(
264        &self,
265        graph: &Graph,
266        output_dir: &Path,
267        files: &mut Vec<String>,
268        generated_at: DateTime<Utc>,
269    ) -> std::io::Result<(HashMap<String, usize>, usize, usize)> {
270        let mut type_counts: HashMap<String, usize> = HashMap::new();
271        let mut anomalous_count = 0;
272        let mut max_feature_dim = 0;
273
274        // Collect all nodes
275        let nodes: Vec<RustGraphNodeOutput> = graph
276            .nodes
277            .values()
278            .map(|node| {
279                let output = self.convert_node(node, generated_at);
280                *type_counts.entry(output.node_type.clone()).or_insert(0) += 1;
281                if output.metadata.is_anomaly {
282                    anomalous_count += 1;
283                }
284                if let Some(ref features) = output.metadata.features {
285                    max_feature_dim = max_feature_dim.max(features.len());
286                }
287                output
288            })
289            .collect();
290
291        // Write nodes
292        let filename = match self.config.output_format {
293            RustGraphOutputFormat::JsonLines => "nodes.jsonl",
294            RustGraphOutputFormat::JsonArray => "nodes.json",
295        };
296        let path = output_dir.join(filename);
297        files.push(filename.to_string());
298
299        let file = File::create(path)?;
300        let mut writer = BufWriter::new(file);
301
302        match self.config.output_format {
303            RustGraphOutputFormat::JsonLines => {
304                for node in &nodes {
305                    serde_json::to_writer(&mut writer, node)?;
306                    writeln!(writer)?;
307                }
308            }
309            RustGraphOutputFormat::JsonArray => {
310                if self.config.pretty_print {
311                    serde_json::to_writer_pretty(&mut writer, &nodes)?;
312                } else {
313                    serde_json::to_writer(&mut writer, &nodes)?;
314                }
315            }
316        }
317
318        writer.flush()?;
319
320        Ok((type_counts, anomalous_count, max_feature_dim))
321    }
322
323    /// Exports edges to file(s).
324    fn export_edges(
325        &self,
326        graph: &Graph,
327        output_dir: &Path,
328        files: &mut Vec<String>,
329        generated_at: DateTime<Utc>,
330    ) -> std::io::Result<(HashMap<String, usize>, usize, usize)> {
331        let mut type_counts: HashMap<String, usize> = HashMap::new();
332        let mut anomalous_count = 0;
333        let mut max_feature_dim = 0;
334
335        // Collect all edges
336        let edges: Vec<RustGraphEdgeOutput> = graph
337            .edges
338            .values()
339            .map(|edge| {
340                let output = self.convert_edge(edge, generated_at);
341                *type_counts.entry(output.edge_type.clone()).or_insert(0) += 1;
342                if output.metadata.is_anomaly {
343                    anomalous_count += 1;
344                }
345                if let Some(ref features) = output.metadata.features {
346                    max_feature_dim = max_feature_dim.max(features.len());
347                }
348                output
349            })
350            .collect();
351
352        // Write edges
353        let filename = match self.config.output_format {
354            RustGraphOutputFormat::JsonLines => "edges.jsonl",
355            RustGraphOutputFormat::JsonArray => "edges.json",
356        };
357        let path = output_dir.join(filename);
358        files.push(filename.to_string());
359
360        let file = File::create(path)?;
361        let mut writer = BufWriter::new(file);
362
363        match self.config.output_format {
364            RustGraphOutputFormat::JsonLines => {
365                for edge in &edges {
366                    serde_json::to_writer(&mut writer, edge)?;
367                    writeln!(writer)?;
368                }
369            }
370            RustGraphOutputFormat::JsonArray => {
371                if self.config.pretty_print {
372                    serde_json::to_writer_pretty(&mut writer, &edges)?;
373                } else {
374                    serde_json::to_writer(&mut writer, &edges)?;
375                }
376            }
377        }
378
379        writer.flush()?;
380
381        Ok((type_counts, anomalous_count, max_feature_dim))
382    }
383
384    /// Converts a GraphNode to RustGraphNodeOutput.
385    pub fn convert_node(
386        &self,
387        node: &GraphNode,
388        generated_at: DateTime<Utc>,
389    ) -> RustGraphNodeOutput {
390        let mut properties = HashMap::new();
391
392        // Add external ID and label as properties
393        properties.insert(
394            "external_id".to_string(),
395            Value::String(node.external_id.clone()),
396        );
397        properties.insert("label".to_string(), Value::String(node.label.clone()));
398
399        // Add node properties
400        if self.config.export_node_properties {
401            for (key, value) in &node.properties {
402                properties.insert(key.clone(), node_property_to_json(value));
403            }
404        }
405
406        // Add categorical features as properties
407        for (key, value) in &node.categorical_features {
408            properties.insert(key.clone(), Value::String(value.clone()));
409        }
410
411        // Build labels map
412        let mut labels = HashMap::new();
413        for (i, label) in node.labels.iter().enumerate() {
414            labels.insert(format!("label_{}", i), label.clone());
415        }
416
417        RustGraphNodeOutput {
418            node_type: node.node_type.as_str().to_string(),
419            id: node.id.to_string(),
420            properties,
421            metadata: RustGraphNodeMetadata {
422                source: self.config.source_name.clone(),
423                generated_at,
424                valid_from: if self.config.include_temporal {
425                    Some(generated_at.naive_utc())
426                } else {
427                    None
428                },
429                valid_to: None,
430                transaction_time: if self.config.include_temporal {
431                    Some(generated_at)
432                } else {
433                    None
434                },
435                labels: if self.config.include_labels {
436                    labels
437                } else {
438                    HashMap::new()
439                },
440                features: if self.config.include_features && !node.features.is_empty() {
441                    Some(node.features.clone())
442                } else {
443                    None
444                },
445                batch_id: self.config.batch_id.clone(),
446                is_anomaly: node.is_anomaly,
447                anomaly_type: node.anomaly_type.clone(),
448            },
449        }
450    }
451
452    /// Converts a GraphEdge to RustGraphEdgeOutput.
453    pub fn convert_edge(
454        &self,
455        edge: &GraphEdge,
456        generated_at: DateTime<Utc>,
457    ) -> RustGraphEdgeOutput {
458        let mut properties = HashMap::new();
459
460        // Add edge properties
461        if self.config.export_edge_properties {
462            for (key, value) in &edge.properties {
463                properties.insert(key.clone(), edge_property_to_json(value));
464            }
465        }
466
467        // Add timestamp if present
468        if let Some(ts) = edge.timestamp {
469            properties.insert("timestamp".to_string(), Value::String(ts.to_string()));
470        }
471
472        // Build labels map
473        let mut labels = HashMap::new();
474        for (i, label) in edge.labels.iter().enumerate() {
475            labels.insert(format!("label_{}", i), label.clone());
476        }
477
478        // Determine valid_from from timestamp
479        let valid_from = if self.config.include_temporal {
480            edge.timestamp
481                .map(|d| d.and_hms_opt(0, 0, 0).unwrap())
482                .or_else(|| Some(generated_at.naive_utc()))
483        } else {
484            None
485        };
486
487        RustGraphEdgeOutput {
488            edge_type: edge.edge_type.as_str().to_string(),
489            id: edge.id.to_string(),
490            source_id: edge.source.to_string(),
491            target_id: edge.target.to_string(),
492            properties,
493            metadata: RustGraphEdgeMetadata {
494                source: self.config.source_name.clone(),
495                generated_at,
496                weight: Some(edge.weight),
497                valid_from,
498                valid_to: None,
499                labels: if self.config.include_labels {
500                    labels
501                } else {
502                    HashMap::new()
503                },
504                features: if self.config.include_features && !edge.features.is_empty() {
505                    Some(edge.features.clone())
506                } else {
507                    None
508                },
509                is_directed: edge.direction == EdgeDirection::Directed,
510                is_anomaly: edge.is_anomaly,
511                anomaly_type: edge.anomaly_type.clone(),
512                batch_id: self.config.batch_id.clone(),
513            },
514        }
515    }
516
517    /// Exports a graph to a writer (for streaming export).
518    pub fn export_to_writer<W: Write>(
519        &self,
520        graph: &Graph,
521        writer: &mut W,
522    ) -> std::io::Result<RustGraphMetadata> {
523        let generated_at = Utc::now();
524        let mut type_counts_nodes: HashMap<String, usize> = HashMap::new();
525        let mut type_counts_edges: HashMap<String, usize> = HashMap::new();
526        let mut anomalous_nodes = 0;
527        let mut anomalous_edges = 0;
528        let mut node_feature_dim = 0;
529        let mut edge_feature_dim = 0;
530
531        // Convert all nodes
532        let nodes: Vec<RustGraphNodeOutput> = graph
533            .nodes
534            .values()
535            .map(|node| {
536                let output = self.convert_node(node, generated_at);
537                *type_counts_nodes
538                    .entry(output.node_type.clone())
539                    .or_insert(0) += 1;
540                if output.metadata.is_anomaly {
541                    anomalous_nodes += 1;
542                }
543                if let Some(ref features) = output.metadata.features {
544                    node_feature_dim = node_feature_dim.max(features.len());
545                }
546                output
547            })
548            .collect();
549
550        // Convert all edges
551        let edges: Vec<RustGraphEdgeOutput> = graph
552            .edges
553            .values()
554            .map(|edge| {
555                let output = self.convert_edge(edge, generated_at);
556                *type_counts_edges
557                    .entry(output.edge_type.clone())
558                    .or_insert(0) += 1;
559                if output.metadata.is_anomaly {
560                    anomalous_edges += 1;
561                }
562                if let Some(ref features) = output.metadata.features {
563                    edge_feature_dim = edge_feature_dim.max(features.len());
564                }
565                output
566            })
567            .collect();
568
569        // Calculate density
570        let n = graph.node_count();
571        let possible_edges = if n > 1 { n * (n - 1) } else { 1 };
572        let graph_density = graph.edge_count() as f64 / possible_edges as f64;
573
574        // Build combined output
575        #[derive(Serialize)]
576        struct CombinedOutput<'a> {
577            nodes: &'a [RustGraphNodeOutput],
578            edges: &'a [RustGraphEdgeOutput],
579        }
580
581        let combined = CombinedOutput {
582            nodes: &nodes,
583            edges: &edges,
584        };
585
586        if self.config.pretty_print {
587            serde_json::to_writer_pretty(writer, &combined)?;
588        } else {
589            serde_json::to_writer(writer, &combined)?;
590        }
591
592        Ok(RustGraphMetadata {
593            name: graph.name.clone(),
594            num_nodes: graph.node_count(),
595            num_edges: graph.edge_count(),
596            node_type_counts: type_counts_nodes,
597            edge_type_counts: type_counts_edges,
598            node_feature_dim,
599            edge_feature_dim,
600            graph_density,
601            anomalous_nodes,
602            anomalous_edges,
603            source: self.config.source_name.clone(),
604            generated_at,
605            output_format: "json".to_string(),
606            files: vec![],
607        })
608    }
609}
610
611/// Converts a NodeProperty to a JSON Value.
612fn node_property_to_json(prop: &NodeProperty) -> Value {
613    match prop {
614        NodeProperty::String(s) => Value::String(s.clone()),
615        NodeProperty::Int(i) => Value::Number((*i).into()),
616        NodeProperty::Float(f) => {
617            serde_json::Number::from_f64(*f).map_or(Value::Null, Value::Number)
618        }
619        NodeProperty::Decimal(d) => Value::String(d.to_string()),
620        NodeProperty::Bool(b) => Value::Bool(*b),
621        NodeProperty::Date(d) => Value::String(d.to_string()),
622        NodeProperty::StringList(v) => {
623            Value::Array(v.iter().map(|s| Value::String(s.clone())).collect())
624        }
625    }
626}
627
628/// Converts an EdgeProperty to a JSON Value.
629fn edge_property_to_json(prop: &EdgeProperty) -> Value {
630    match prop {
631        EdgeProperty::String(s) => Value::String(s.clone()),
632        EdgeProperty::Int(i) => Value::Number((*i).into()),
633        EdgeProperty::Float(f) => {
634            serde_json::Number::from_f64(*f).map_or(Value::Null, Value::Number)
635        }
636        EdgeProperty::Decimal(d) => Value::String(d.to_string()),
637        EdgeProperty::Bool(b) => Value::Bool(*b),
638        EdgeProperty::Date(d) => Value::String(d.to_string()),
639    }
640}
641
642#[cfg(test)]
643mod tests {
644    use super::*;
645    use crate::models::{EdgeType, GraphType, NodeType};
646    use tempfile::tempdir;
647
648    fn create_test_graph() -> Graph {
649        let mut graph = Graph::new("test_graph", GraphType::Transaction);
650
651        let n1 = graph.add_node(
652            GraphNode::new(0, NodeType::Account, "1000".to_string(), "Cash".to_string())
653                .with_feature(0.5)
654                .with_feature(0.3)
655                .with_categorical("account_type", "Asset"),
656        );
657        let n2 = graph.add_node(
658            GraphNode::new(0, NodeType::Account, "2000".to_string(), "AP".to_string())
659                .with_feature(0.8)
660                .with_feature(0.2)
661                .as_anomaly("unusual_balance"),
662        );
663        let n3 = graph.add_node(
664            GraphNode::new(
665                0,
666                NodeType::Vendor,
667                "V001".to_string(),
668                "Acme Corp".to_string(),
669            )
670            .with_feature(1.0),
671        );
672
673        graph.add_edge(
674            GraphEdge::new(0, n1, n2, EdgeType::Transaction)
675                .with_weight(1000.0)
676                .with_feature(6.9)
677                .with_timestamp(chrono::NaiveDate::from_ymd_opt(2024, 1, 15).unwrap()),
678        );
679        graph.add_edge(
680            GraphEdge::new(0, n2, n3, EdgeType::Transaction)
681                .with_weight(500.0)
682                .as_anomaly("split_transaction"),
683        );
684
685        graph.compute_statistics();
686        graph
687    }
688
689    #[test]
690    fn test_rustgraph_export_jsonl() {
691        let graph = create_test_graph();
692        let dir = tempdir().unwrap();
693
694        let config = RustGraphExportConfig {
695            output_format: RustGraphOutputFormat::JsonLines,
696            ..Default::default()
697        };
698        let exporter = RustGraphExporter::new(config);
699        let metadata = exporter.export(&graph, dir.path()).unwrap();
700
701        assert_eq!(metadata.num_nodes, 3);
702        assert_eq!(metadata.num_edges, 2);
703        assert_eq!(metadata.anomalous_nodes, 1);
704        assert_eq!(metadata.anomalous_edges, 1);
705        assert!(dir.path().join("nodes.jsonl").exists());
706        assert!(dir.path().join("edges.jsonl").exists());
707        assert!(dir.path().join("metadata.json").exists());
708    }
709
710    #[test]
711    fn test_rustgraph_export_json_array() {
712        let graph = create_test_graph();
713        let dir = tempdir().unwrap();
714
715        let config = RustGraphExportConfig {
716            output_format: RustGraphOutputFormat::JsonArray,
717            pretty_print: true,
718            ..Default::default()
719        };
720        let exporter = RustGraphExporter::new(config);
721        let metadata = exporter.export(&graph, dir.path()).unwrap();
722
723        assert_eq!(metadata.num_nodes, 3);
724        assert!(dir.path().join("nodes.json").exists());
725        assert!(dir.path().join("edges.json").exists());
726    }
727
728    #[test]
729    fn test_convert_node() {
730        let node = GraphNode::new(
731            42,
732            NodeType::Vendor,
733            "V001".to_string(),
734            "Test Vendor".to_string(),
735        )
736        .with_feature(1.0)
737        .with_feature(2.0)
738        .with_categorical("region", "US")
739        .with_label("high_risk");
740
741        let config = RustGraphExportConfig::default();
742        let exporter = RustGraphExporter::new(config);
743        let output = exporter.convert_node(&node, Utc::now());
744
745        assert_eq!(output.node_type, "Vendor");
746        assert_eq!(output.id, "42");
747        assert_eq!(
748            output.properties.get("external_id"),
749            Some(&Value::String("V001".to_string()))
750        );
751        assert_eq!(
752            output.properties.get("region"),
753            Some(&Value::String("US".to_string()))
754        );
755        assert_eq!(output.metadata.features, Some(vec![1.0, 2.0]));
756        assert!(!output.metadata.is_anomaly);
757    }
758
759    #[test]
760    fn test_convert_edge() {
761        let edge = GraphEdge::new(99, 1, 2, EdgeType::Approval)
762            .with_weight(5000.0)
763            .with_feature(0.5)
764            .with_timestamp(chrono::NaiveDate::from_ymd_opt(2024, 6, 15).unwrap())
765            .as_anomaly("threshold_breach");
766
767        let config = RustGraphExportConfig::default();
768        let exporter = RustGraphExporter::new(config);
769        let output = exporter.convert_edge(&edge, Utc::now());
770
771        assert_eq!(output.edge_type, "Approval");
772        assert_eq!(output.id, "99");
773        assert_eq!(output.source_id, "1");
774        assert_eq!(output.target_id, "2");
775        assert_eq!(output.metadata.weight, Some(5000.0));
776        assert!(output.metadata.is_directed);
777        assert!(output.metadata.is_anomaly);
778        assert_eq!(
779            output.metadata.anomaly_type,
780            Some("threshold_breach".to_string())
781        );
782    }
783
784    #[test]
785    fn test_export_to_writer() {
786        let graph = create_test_graph();
787        let mut buffer = Vec::new();
788
789        let config = RustGraphExportConfig {
790            include_features: true,
791            include_temporal: true,
792            ..Default::default()
793        };
794        let exporter = RustGraphExporter::new(config);
795        let metadata = exporter.export_to_writer(&graph, &mut buffer).unwrap();
796
797        assert_eq!(metadata.num_nodes, 3);
798        assert_eq!(metadata.num_edges, 2);
799
800        let output: serde_json::Value = serde_json::from_slice(&buffer).unwrap();
801        assert!(output.get("nodes").is_some());
802        assert!(output.get("edges").is_some());
803    }
804
805    #[test]
806    fn test_node_type_counts() {
807        let graph = create_test_graph();
808        let dir = tempdir().unwrap();
809
810        let exporter = RustGraphExporter::new(RustGraphExportConfig::default());
811        let metadata = exporter.export(&graph, dir.path()).unwrap();
812
813        assert_eq!(metadata.node_type_counts.get("Account"), Some(&2));
814        assert_eq!(metadata.node_type_counts.get("Vendor"), Some(&1));
815    }
816
817    #[test]
818    fn test_without_features() {
819        let node = GraphNode::new(1, NodeType::Account, "1000".to_string(), "Cash".to_string())
820            .with_feature(1.0);
821
822        let config = RustGraphExportConfig {
823            include_features: false,
824            ..Default::default()
825        };
826        let exporter = RustGraphExporter::new(config);
827        let output = exporter.convert_node(&node, Utc::now());
828
829        assert!(output.metadata.features.is_none());
830    }
831
832    #[test]
833    fn test_with_batch_id() {
834        let node = GraphNode::new(1, NodeType::Account, "1000".to_string(), "Cash".to_string());
835
836        let config = RustGraphExportConfig {
837            batch_id: Some("batch_001".to_string()),
838            ..Default::default()
839        };
840        let exporter = RustGraphExporter::new(config);
841        let output = exporter.convert_node(&node, Utc::now());
842
843        assert_eq!(output.metadata.batch_id, Some("batch_001".to_string()));
844    }
845}
datasynth_graph/exporters/rustgraph.rs

datasynth_graph/exporters/
rustgraph.rs
