Added BayesNet_test

core/BayesNet.py

@@ -47,8 +47,13 @@ class BayesNet(object):
             raise Exception("There is no node with name "+node_name+" in the bayesnet")        
 
     def get_nodes(self, node_names):
-        for node_name in node_names:
-            yield self.get_node(node_name)
+        nodes = []
+        if not node_names:
+            nodes = self.bn.nodes()
+        else:
+            for node_name in node_names:
+                nodes.append(self.get_node(node_name))
+        return nodes
 
     def get_parents(self, node):
         raise Exception("Called unimplemented function")

core/tests/BayesNet_test.py

+import unittest
+
+from core.BayesNet import BayesNet
+from core.Node import Node
+
+class NodeAddAndRemoveTestCase(unittest.TestCase):
+    def setUp(self):
+        self.bn = BayesNet()
+
+    def tearDown(self):
+        self.bn = None
+
+    def test_add(self):
+        n = Node("Some Node")
+        self.bn.add_node(n)
+        self.assertEqual(n,self.bn.get_node("Some Node"))
+        self.assertTrue(n in self.bn.get_nodes(["Some Node"]))
+
+        node_with_same_name=Node("Some Node")
+        self.assertRaises(Exception, self.bn.add_node, node_with_same_name)
+
+    def test_remove(self):
+        n = Node("Some Node to remove")
+        self.bn.add_node(n)
+        self.bn.remove_node(n)
+        self.assertFalse(n in self.bn.get_nodes())
+        
+
+
+#include this so you can run this test without nose
+if __name__ == '__main__':
+    unittest.main()