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()

[((44, 4, 44, 19), 'unittest.main', 'unittest.main', ({}, {}), '()', False, 'import unittest\n'), ((24, 22, 24, 43), 'molecule.parse_molecule', 'parse_molecule', ({(24, 37, 24, 42): '"""H2O"""'}, {}), "('H2O')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((29, 22, 29, 47), 'molecule.parse_molecule', 'parse_molecule', ({(29, 37, 29, 46): '"""Mg(OH)2"""'}, {}), "('Mg(OH)2')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((35, 22, 35, 53), 'molecule.parse_molecule', 'parse_molecule', ({(35, 37, 35, 52): '"""K4[ON(SO3)2]2"""'}, {}), "('K4[ON(SO3)2]2')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((8, 25, 8, 43), 'molecule.onize_formula', 'onize_formula', ({(8, 39, 8, 42): '"""H"""'}, {}), "('H')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((9, 25, 9, 45), 'molecule.onize_formula', 'onize_formula', ({(9, 39, 9, 44): '"""H2O"""'}, {}), "('H2O')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((10, 25, 10, 49), 'molecule.onize_formula', 'onize_formula', ({(10, 39, 10, 48): '"""Mg(OH)2"""'}, {}), "('Mg(OH)2')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((11, 25, 11, 55), 'molecule.onize_formula', 'onize_formula', ({(11, 39, 11, 54): '"""K4[ON(SO3)2]2"""'}, {}), "('K4[ON(SO3)2]2')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((14, 25, 14, 67), 'molecule.update_equation_with_multiplier', 'update_equation_with_multiplier', ({(14, 57, 14, 61): '"""H1"""', (14, 63, 14, 66): '"""2"""'}, {}), "('H1', '2')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((15, 25, 15, 80), 'molecule.update_equation_with_multiplier', 'update_equation_with_multiplier', ({(15, 57, 15, 74): '"""K4[O1N1(SO3)2]2"""', (15, 76, 15, 79): '"""2"""'}, {}), "('K4[O1N1(SO3)2]2', '2')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((18, 25, 18, 46), 'molecule.flaten_formula', 'flaten_formula', ({(18, 40, 18, 45): '"""H2O"""'}, {}), "('H2O')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((19, 25, 19, 49), 'molecule.flaten_formula', 'flaten_formula', ({(19, 40, 19, 48): '"""[H1]2O"""'}, {}), "('[H1]2O')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((20, 25, 20, 54), 'molecule.flaten_formula', 'flaten_formula', ({(20, 40, 20, 53): '"""M1g1(O1H1)2"""'}, {}), "('M1g1(O1H1)2')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n'), ((21, 25, 21, 59), 'molecule.flaten_formula', 'flaten_formula', ({(21, 40, 21, 58): '"""K4[O1N1(S1O3)2]2"""'}, {}), "('K4[O1N1(S1O3)2]2')", False, 'from molecule import onize_formula, update_equation_with_multiplier, flaten_formula, parse_molecule\n')]