import unittest

from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule

class MoleculeParserTestCases(unittest.TestCase):

    def test_onizing_formulas(self):
        self.assertEqual(onize_formula('H'), 'H1')
        self.assertEqual(onize_formula('H2O'), 'H2O1')
        self.assertEqual(onize_formula('Mg(OH)2'), 'Mg1(O1H1)2')
        self.assertEqual(onize_formula('K4[ON(SO3)2]2'), 'K4[O1N1(S1O3)2]2')

    def test_updating_formula_with_multipler(self):
        self.assertEqual(update_equation_with_multiplier('H1', '2'), 'H2')
        self.assertEqual(update_equation_with_multiplier('K4[O1N1(SO3)2]2', '2'), 'K8[O2N2(SO6)4]4')

    def test_flatting_formula(self):
        self.assertEqual(flaten_formula('H2O'), 'H2O')
        self.assertEqual(flaten_formula('[H1]2O'), 'H2O')
        self.assertEqual(flaten_formula('M1g1(O1H1)2'), 'M1g1O2H2')
        self.assertEqual(flaten_formula('K4[O1N1(S1O3)2]2'), 'K4O2N2S4O12')

    def test_full_parsing(self):
        parsed_mole = parse_molecule('H2O')
        self.assertEqual(len(parsed_mole.keys()), 2)
        self.assertEqual(parsed_mole['H'], 2)
        self.assertEqual(parsed_mole['O'], 1)

        parsed_mole = parse_molecule('Mg(OH)2')
        self.assertEqual(len(parsed_mole.keys()), 3)
        self.assertEqual(parsed_mole['H'], 2)
        self.assertEqual(parsed_mole['O'], 2)
        self.assertEqual(parsed_mole['Mg'], 1)

        parsed_mole = parse_molecule('K4[ON(SO3)2]2')
        self.assertEqual(len(parsed_mole.keys()), 4)
        self.assertEqual(parsed_mole['K'], 4)
        self.assertEqual(parsed_mole['O'], 14)
        self.assertEqual(parsed_mole['N'], 2)
        self.assertEqual(parsed_mole['S'], 4)


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

[((1768, 1783), 'unittest.main', 'unittest.main', ([], {}), '()\n', (1781, 1783), False, 'import unittest\n'), ((1034, 1055), 'molecule.parse_molecule', 'parse_molecule', (['"""H2O"""'], {}), "('H2O')\n", (1048, 1055), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((1224, 1249), 'molecule.parse_molecule', 'parse_molecule', (['"""Mg(OH)2"""'], {}), "('Mg(OH)2')\n", (1238, 1249), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((1465, 1496), 'molecule.parse_molecule', 'parse_molecule', (['"""K4[ON(SO3)2]2"""'], {}), "('K4[ON(SO3)2]2')\n", (1479, 1496), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((231, 249), 'molecule.onize_formula', 'onize_formula', (['"""H"""'], {}), "('H')\n", (244, 249), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((282, 302), 'molecule.onize_formula', 'onize_formula', (['"""H2O"""'], {}), "('H2O')\n", (295, 302), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((337, 361), 'molecule.onize_formula', 'onize_formula', (['"""Mg(OH)2"""'], {}), "('Mg(OH)2')\n", (350, 361), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((402, 432), 'molecule.onize_formula', 'onize_formula', (['"""K4[ON(SO3)2]2"""'], {}), "('K4[ON(SO3)2]2')\n", (415, 432), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((532, 574), 'molecule.update_equation_with_multiplier', 'update_equation_with_multiplier', (['"""H1"""', '"""2"""'], {}), "('H1', '2')\n", (563, 574), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((607, 662), 'molecule.update_equation_with_multiplier', 'update_equation_with_multiplier', (['"""K4[O1N1(SO3)2]2"""', '"""2"""'], {}), "('K4[O1N1(SO3)2]2', '2')\n", (638, 662), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((746, 767), 'molecule.flaten_formula', 'flaten_formula', (['"""H2O"""'], {}), "('H2O')\n", (760, 767), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((801, 825), 'molecule.flaten_formula', 'flaten_formula', (['"""[H1]2O"""'], {}), "('[H1]2O')\n", (815, 825), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((859, 888), 'molecule.flaten_formula', 'flaten_formula', (['"""M1g1(O1H1)2"""'], {}), "('M1g1(O1H1)2')\n", (873, 888), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((927, 961), 'molecule.flaten_formula', 'flaten_formula', (['"""K4[O1N1(S1O3)2]2"""'], {}), "('K4[O1N1(S1O3)2]2')\n", (941, 961), False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n')]