From 5444c619a2ef402af095cca376bd47c1db8f8120 Mon Sep 17 00:00:00 2001
From: AlessandroBregoli <alessandroxciv@gmail.com>
Date: Thu, 14 Apr 2022 14:22:53 +0200
Subject: [PATCH] Added tests

---
 tests/structure_learning.rs | 40 +++++++++++++++++++++++++++++++------
 tests/tools.rs              | 24 +++++++++++++++++++++-
 2 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/tests/structure_learning.rs b/tests/structure_learning.rs
index ad18c18..d24170a 100644
--- a/tests/structure_learning.rs
+++ b/tests/structure_learning.rs
@@ -105,8 +105,7 @@ pub fn learn_ternary_net_2_nodes_hill_climbing_bic() {
 }
 
 
-
-fn learn_mixed_discrete_net_3_nodes<T: StructureLearningAlgorithm> (sl: T) {
+fn get_mixed_discrete_net_3_nodes_with_data() -> (CtbnNetwork, Dataset) {
     let mut net = CtbnNetwork::init();
     let n1 = net
         .add_node(generate_discrete_time_continous_node(String::from("n1"),3))
@@ -161,11 +160,15 @@ fn learn_mixed_discrete_net_3_nodes<T: StructureLearningAlgorithm> (sl: T) {
 
 
     let data = trajectory_generator(&net, 300, 30.0, Some(6347747169756259),);
-    let net = sl.fit_transform(net, &data);
+    return (net, data);
+}
 
-    assert_eq!(BTreeSet::new(), net.get_parent_set(n1));
-    assert_eq!(BTreeSet::from_iter(vec![n1]), net.get_parent_set(n2));
-    assert_eq!(BTreeSet::from_iter(vec![n1, n2]), net.get_parent_set(n3));
+fn learn_mixed_discrete_net_3_nodes<T: StructureLearningAlgorithm> (sl: T) {
+    let (net, data) = get_mixed_discrete_net_3_nodes_with_data();
+    let net = sl.fit_transform(net, &data);
+    assert_eq!(BTreeSet::new(), net.get_parent_set(0));
+    assert_eq!(BTreeSet::from_iter(vec![0]), net.get_parent_set(1));
+    assert_eq!(BTreeSet::from_iter(vec![0, 1]), net.get_parent_set(2));
 }
 
 
@@ -182,3 +185,28 @@ pub fn learn_mixed_discrete_net_3_nodes_hill_climbing_bic() {
     let hl = HillClimbing::init(bic, None);
     learn_mixed_discrete_net_3_nodes(hl);
 }
+
+
+
+fn learn_mixed_discrete_net_3_nodes_1_parent_constraint<T: StructureLearningAlgorithm> (sl: T) {
+    let (net, data) = get_mixed_discrete_net_3_nodes_with_data();
+    let net = sl.fit_transform(net, &data);
+    assert_eq!(BTreeSet::new(), net.get_parent_set(0));
+    assert_eq!(BTreeSet::from_iter(vec![0]), net.get_parent_set(1));
+    assert_eq!(BTreeSet::from_iter(vec![0]), net.get_parent_set(2));
+}
+
+
+#[test]
+pub fn learn_mixed_discrete_net_3_nodes_hill_climbing_ll_1_parent_constraint() {
+    let ll = LogLikelihood::init(1, 1.0);
+    let hl = HillClimbing::init(ll, Some(1));
+    learn_mixed_discrete_net_3_nodes_1_parent_constraint(hl);
+}
+
+#[test]
+pub fn learn_mixed_discrete_net_3_nodes_hill_climbing_bic_1_parent_constraint() {
+    let bic = BIC::init(1, 1.0);
+    let hl = HillClimbing::init(bic, Some(1));
+    learn_mixed_discrete_net_3_nodes_1_parent_constraint(hl);
+}
diff --git a/tests/tools.rs b/tests/tools.rs
index 28b3e0d..c341b8c 100644
--- a/tests/tools.rs
+++ b/tests/tools.rs
@@ -5,7 +5,7 @@ use rustyCTBN::ctbn::*;
 use rustyCTBN::node;
 use rustyCTBN::params;
 use std::collections::BTreeSet;
-use ndarray::arr3;
+use ndarray::{arr1, arr2, arr3};
 
 
 
@@ -43,3 +43,25 @@ fn run_sampling() {
 }
 
 
+#[test]
+#[should_panic]
+    fn trajectory_wrong_shape() {
+    let time = arr1(&[0.0, 0.2]);
+    let events = arr2(&[[0,3]]);
+    Trajectory::init(time, events);
+}
+
+
+#[test]
+#[should_panic]
+fn dataset_wrong_shape() {
+    let time = arr1(&[0.0, 0.2]);
+    let events = arr2(&[[0,3], [1,2]]);
+    let t1 = Trajectory::init(time, events);
+
+
+    let time = arr1(&[0.0, 0.2]);
+    let events = arr2(&[[0,3,3], [1,2,3]]);
+    let t2 = Trajectory::init(time, events);
+    Dataset::init(vec![t1, t2]);
+}