#!/usr/bin/python
from PIPRes.tree import *
from PIPRes.wrap.external_program import *
import unittest, re, cStringIO

class TestExternalProgram(unittest.TestCase):
    def genericMapLeaves(self, tr):
        ls = tr.getLeafSet()
        mappedTree, taxa, rows, toExternalTranslation = mapTreeToLeafSubset(tr)
        self.assertEqual(taxa, [('tax%d'% (i + 1)) for i in rows])
        self.assertEqual(ls, [toExternalTranslation(i) + 1 for i in range(len(ls))])
        return mappedTree, taxa, rows, toExternalTranslation

    def testMapLeaves(self):
        tr = numberedLeafTree('(5,1,(7,4))')
        mappedTree, taxa, rows, toExternalTranslation = self.genericMapLeaves(tr)
        self.assertEqual(str(mappedTree), '(3,1,(4,2))')
        self.assertEqual(taxa, ['tax1', 'tax4', 'tax5', 'tax7'])
        self.assertEqual(rows, [0, 3, 4, 6])

if __name__ == '__main__':
    unittest.main()