from fragutils.utils.vector_utils import get_exit_vector_for_xe_smi
import unittest


class VectorTest(unittest.TestCase):
    def test_xe_smi(self):
        # Some simple test data - these should all be zero
        input_data = ["[100Xe]c1ccc([101Xe])cc1", "[101Xe]c1ccccc1"]
        output_data = [
            {
                100: 0.0,
                101: 0.0,
            },
            {101: 0.0},
        ]
        # Now iterate over
        for i, smi in enumerate(input_data):
            ans = get_exit_vector_for_xe_smi(smi)
            test_dict = output_data[i]
            for key in test_dict:
                self.assertAlmostEqual(ans[key], test_dict[key], 2)
