use arcweight::prelude::*;
use arcweight::utils::{draw_fst, draw_fst_default, DrawingConfig};
#[test]
fn test_draw_fst_default() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
let s1 = fst.add_state();
fst.set_start(s0);
fst.set_final(s1, TropicalWeight::new(0.5));
fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.0), s1));
let dot = draw_fst_default(&fst).unwrap();
assert!(dot.contains("digraph FST"));
assert!(dot.contains("rankdir=LR"));
assert!(dot.contains("0 -> 1"));
}
#[test]
fn test_draw_fst_with_config() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
let s1 = fst.add_state();
fst.set_start(s0);
fst.set_final(s1, TropicalWeight::one());
fst.add_arc(s0, Arc::new(1, 2, TropicalWeight::new(1.5), s1));
let config = DrawingConfig {
show_weights: true,
show_state_ids: true,
show_final_weights: true,
use_symbols: false,
horizontal: true,
node_shape: "box".to_string(),
node_color: "white".to_string(),
start_color: "blue".to_string(),
final_color: "green".to_string(),
};
let dot = draw_fst(&fst, config, None, None).unwrap();
assert!(dot.contains("digraph FST"));
assert!(dot.contains("shape=box"));
assert!(dot.contains("fillcolor=blue")); assert!(dot.contains("fillcolor=green")); assert!(dot.contains("1:2 / 1.5")); }
#[test]
fn test_draw_fst_with_symbols() {
let mut fst = VectorFst::<TropicalWeight>::new();
let mut symbols = SymbolTable::new();
let hello_id = symbols.add_symbol("hello");
let world_id = symbols.add_symbol("world");
let s0 = fst.add_state();
let s1 = fst.add_state();
fst.set_start(s0);
fst.set_final(s1, TropicalWeight::one());
fst.add_arc(s0, Arc::new(hello_id, world_id, TropicalWeight::one(), s1));
let config = DrawingConfig {
use_symbols: true,
..Default::default()
};
let dot = draw_fst(&fst, config, Some(&symbols), Some(&symbols)).unwrap();
assert!(dot.contains("hello:world"));
}
#[test]
fn test_draw_fst_no_weights() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
let s1 = fst.add_state();
fst.set_start(s0);
fst.set_final(s1, TropicalWeight::one());
fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::new(1.5), s1));
let config = DrawingConfig {
show_weights: false,
..Default::default()
};
let dot = draw_fst(&fst, config, None, None).unwrap();
assert!(!dot.contains("/ 1.5")); }
#[test]
fn test_draw_fst_vertical_layout() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
let s1 = fst.add_state();
fst.set_start(s0);
fst.set_final(s1, TropicalWeight::one());
fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
let config = DrawingConfig {
horizontal: false,
..Default::default()
};
let dot = draw_fst(&fst, config, None, None).unwrap();
assert!(!dot.contains("rankdir=LR")); }
#[test]
fn test_draw_fst_multiple_states() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
let s1 = fst.add_state();
let s2 = fst.add_state();
fst.set_start(s0);
fst.set_final(s2, TropicalWeight::one());
fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s2));
let dot = draw_fst_default(&fst).unwrap();
assert!(dot.contains("0 -> 1"));
assert!(dot.contains("1 -> 2"));
assert!(dot.matches("0").count() >= 2); }
#[test]
fn test_draw_fst_epsilon_arcs() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
let s1 = fst.add_state();
fst.set_start(s0);
fst.set_final(s1, TropicalWeight::one());
fst.add_arc(s0, Arc::new(0, 0, TropicalWeight::one(), s1));
let dot = draw_fst_default(&fst).unwrap();
assert!(dot.contains("0 -> 1"));
}
#[test]
fn test_draw_fst_empty() {
let fst = VectorFst::<TropicalWeight>::new();
let dot = draw_fst_default(&fst).unwrap();
assert!(dot.contains("digraph FST"));
assert!(dot.contains("}"));
}
#[test]
fn test_draw_fst_start_and_final_same() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
fst.set_start(s0);
fst.set_final(s0, TropicalWeight::one());
let dot = draw_fst_default(&fst).unwrap();
assert!(dot.contains("0"));
}
#[test]
fn test_draw_fst_complex_structure() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
let s1 = fst.add_state();
let s2 = fst.add_state();
let s3 = fst.add_state();
fst.set_start(s0);
fst.set_final(s2, TropicalWeight::new(0.5));
fst.set_final(s3, TropicalWeight::new(0.3));
fst.add_arc(s0, Arc::new(1, 1, TropicalWeight::one(), s1));
fst.add_arc(s1, Arc::new(2, 2, TropicalWeight::one(), s2));
fst.add_arc(s0, Arc::new(3, 3, TropicalWeight::one(), s3));
let dot = draw_fst_default(&fst).unwrap();
assert!(dot.contains("0 -> 1"));
assert!(dot.contains("1 -> 2"));
assert!(dot.contains("0 -> 3"));
}
#[test]
fn test_draw_fst_custom_node_shape() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
fst.set_start(s0);
fst.set_final(s0, TropicalWeight::one());
let config = DrawingConfig {
node_shape: "diamond".to_string(),
..Default::default()
};
let dot = draw_fst(&fst, config, None, None).unwrap();
assert!(dot.contains("shape=diamond"));
}
#[test]
fn test_draw_fst_final_weight_display() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
let s1 = fst.add_state();
fst.set_start(s0);
fst.set_final(s1, TropicalWeight::new(0.75));
let config = DrawingConfig {
show_final_weights: true,
..Default::default()
};
let dot = draw_fst(&fst, config, None, None).unwrap();
assert!(dot.contains("0.75")); }
#[test]
fn test_draw_fst_no_final_weight_display() {
let mut fst = VectorFst::<TropicalWeight>::new();
let s0 = fst.add_state();
let s1 = fst.add_state();
fst.set_start(s0);
fst.set_final(s1, TropicalWeight::new(0.75));
let config = DrawingConfig {
show_final_weights: false,
..Default::default()
};
let dot = draw_fst(&fst, config, None, None).unwrap();
assert!(dot.contains("digraph FST"));
}