1 files changed, 118 insertions, 0 deletions
diff --git a/backend/tests/test_gen_otol_data.py b/backend/tests/test_gen_otol_data.py
new file mode 100644
index 0000000..25e65e3
--- /dev/null
+++ b/backend/tests/test_gen_otol_data.py
@@ -0,0 +1,118 @@
+import unittest
+import tempfile, os
+
+from tests.common import createTestFile, readTestDbTable
+from tol_data.gen_otol_data import genData
+
+def runGenData(treeFileContents: str, annFileContents: str, pickedFileContents: str):
+	""" Sets up files to be read by genData(), runs it, reads the output database, and returns node+edge info """
+	with tempfile.TemporaryDirectory() as tempDir:
+		# Create temp tree file
+		treeFile = os.path.join(tempDir, 'tree.tre')
+		createTestFile(treeFile, treeFileContents)
+		# Create temp annotations file
+		annFile = os.path.join(tempDir, 'ann.json')
+		createTestFile(annFile, annFileContents)
+		# Create temp picked names file
+		pickedFile = os.path.join(tempDir, 'pn.txt')
+		createTestFile(pickedFile, pickedFileContents)
+		# Run genData()
+		dbFile = os.path.join(tempDir, 'data.db')
+		genData(treeFile, annFile, pickedFile, dbFile)
+		# Read database
+		nodes = readTestDbTable(dbFile, 'SELECT name, id, tips FROM nodes')
+		edges = readTestDbTable(dbFile, 'SELECT parent, child, p_support FROM edges')
+	return nodes, edges
+
+class TestGenData(unittest.TestCase):
+	def setUp(self):
+		self.maxDiff = None # Remove output-diff size limit
+	def test_newick(self):
+		treeFileContents = """
+			(
+				'land plants ott2',
+				(
+					'TRAVELLER''s tree ott100',
+					(domestic_banana_ott4, (lemon_ott6, orange_ott7)citrus_ott5)mrcaott4ott5
+				) mrcaott100ott4,
+				'Highly  Unu2u8| name!!  ott999',
+				'citrus ott230'
+			)cellular_organisms_ott1;"""
+		annFileContents = '{"nodes": {}}'
+		pickedFileContents = ''
+		nodes, edges = runGenData(treeFileContents, annFileContents, pickedFileContents)
+		self.assertEqual(nodes, {
+			('land plants', 'ott2', 1),
+			('traveller\'s tree', 'ott100', 1),
+			('domestic banana', 'ott4', 1),
+			('lemon', 'ott6', 1),
+			('orange', 'ott7', 1),
+			('citrus', 'ott5', 2),
+			('[citrus + domestic banana]', 'mrcaott4ott5', 3),
+			('[citrus + traveller\'s tree]', 'mrcaott100ott4', 4),
+			('highly  unu2u8| name!! ', 'ott999', 1),
+			('citrus [2]', 'ott230', 1),
+			('cellular organisms', 'ott1', 7),
+		})
+		self.assertEqual(edges, {
+			('cellular organisms', 'land plants', 0),
+			('cellular organisms', '[citrus + traveller\'s tree]', 0),
+			('cellular organisms', 'highly  unu2u8| name!! ', 0),
+			('cellular organisms', 'citrus [2]', 0),
+			('[citrus + traveller\'s tree]', 'traveller\'s tree', 0),
+			('[citrus + traveller\'s tree]', '[citrus + domestic banana]', 0),
+			('[citrus + domestic banana]', 'domestic banana', 0),
+			('[citrus + domestic banana]', 'citrus', 0),
+			('citrus', 'lemon', 0),
+			('citrus', 'orange', 0),
+		})
+	def test_newick_invalid(self):
+		with self.assertRaises(Exception):
+			runGenData('(A,B,(C,D));', '{"nodes": {}}', '')
+	def test_annotations(self):
+		treeFileContents = '(two_ott2, three_ott3, four_ott4)one_ott1;'
+		annFileContents = """
+			{
+				"date_completed": "xxx",
+				"nodes": {
+					"ott3": {
+						"supported_by": {
+							"tree1": "node1"
+						}
+					},
+					"ott4": {
+						"supported_by": {
+							"tree1": "node2",
+							"tree2": "node100"
+						},
+						"conflicts_with": {
+							"tree3": ["x", "y"]
+						}
+					}
+				}
+			}"""
+		nodes, edges = runGenData(treeFileContents, annFileContents, '')
+		self.assertEqual(nodes, {
+			('one', 'ott1', 3),
+			('two', 'ott2', 1),
+			('three', 'ott3', 1),
+			('four', 'ott4', 1),
+		})
+		self.assertEqual(edges, {
+			('one', 'two', 0),
+			('one', 'three', 1),
+			('one', 'four', 0),
+		})
+	def test_picked_names_file(self):
+		treeFileContents = '(one_ott2, two_ott3)one_ott1;'
+		pickedFileContents = 'one|ott2'
+		nodes, edges = runGenData(treeFileContents, '{"nodes": {}}', pickedFileContents)
+		self.assertEqual(nodes, {
+			('one [2]', 'ott1', 2),
+			('one', 'ott2', 1),
+			('two', 'ott3', 1),
+		})
+		self.assertEqual(edges, {
+			('one [2]', 'one', 0),
+			('one [2]', 'two', 0),
+		})