add tests

tests/test_cpi.py

@@ -5,7 +5,7 @@
 
 from rail import CPI
 
-class TestTree(unittest.TestCase):
+class TestCPI(unittest.TestCase):
     """
     Class to test a CPI.
     """

tests/test_threat_sources.py

@@ -0,0 +1,24 @@
+"""
+Tests for the ThreatSource and ThreatSources classes
+"""
+import unittest
+
+from rail import ThreatSources
+
+class TestThreatSource(unittest.TestCase):
+    """
+    Class to test a threat source
+    """
+    def setUp(self):
+        self.threat_sources = ThreatSources()
+        self.threat_sources.new('test')
+
+    def test_threat_sources(self):
+        """
+        Test threat sources
+        """
+        self.assertEqual(self.threat_sources['test']['name'], 'test')
+
+
+if __name__ == '__main__':
+    unittest.main()

tests/test_tree.py

@@ -39,7 +39,10 @@ def test_print(self):
             mock_print.reset_mock()
             self.assertEqual(self.tree.to_print(), None)
             mock_print.assert_called()
-            mock_print.assert_has_calls([mock.call('/test tree'), mock.call('/test tree/test child')])
+            mock_print.assert_has_calls([
+                mock.call('/test tree'),
+                mock.call('/test tree/test child')
+            ])
 
     def test_to_latex(self):
         """
@@ -49,11 +52,11 @@ def test_to_latex(self):
             self.assertEqual(self.tree.to_latex(), None)
             mock_print.assert_called()
             mock_print.assert_has_calls([
-                mock.call('child { node{test tree}'), 
+                mock.call('child { node{test tree}'),
                 mock.call('child { node{test child}'),
                 mock.call('}'),
-                mock.call('}')]
-            )   
+                mock.call('}')
+            ])
 
     def test_to_dict_list(self):
         """