Пример #1
0
class SEGYSortingTestCase(unittest.TestCase):
    """
    Test cases for the sorting module. 
    """

    def setUp(self):
        """
        Create a SEGYFile object to sort.
        """
        self.sgy = SEGYFile()
        self.attr1 = "ensemble_number"
        self.attr2 = "trace_number_within_the_ensemble"
        self.val1 = [2, 1, 2, 1, 1, 2]
        self.val2 = [5, 4, 3, 2, 1, 0]
        for i, k1 in enumerate(self.val1):
            k2 = self.val2[i]
            self.sgy.traces.append(SEGYTrace())
            self.sgy.traces[i].header.__setattr__(self.attr1, k1)
            self.sgy.traces[i].header.__setattr__(self.attr2, k2)

    def test_decorate_undecorate(self):
        """
        Decorate should be reversed by undecorate. 
        """
        # Decorated traces should have attribute values prepended.
        self.sgy._decorate_traces(self.attr1, self.attr2)
        for i, tr in enumerate(self.sgy.traces):
            self.assertEqual([self.val1[i], self.val2[i], i, tr[-1]], self.sgy.traces[i])
        # Undecorating decorated traces should return the original traces
        self.sgy._undecorate_traces()
        for i, tr in enumerate(self.sgy.traces):
            self.assertTrue(isinstance(tr, SEGYTrace))

    def test_sort_traces(self):
        """
        sort_traces should sort traces in place.
        """
        # sort the traces
        self.sgy.sort_traces(self.attr1, self.attr2)
        # traces should now be sorted
        # check 1st key
        self.assertEqual([tr.header.__getattribute__(self.attr1) for tr in self.sgy.traces], sorted(self.val1))
        # check all keys
        values = []
        for i, row in enumerate(zip(self.val1, self.val2)):
            values.append([v for v in row])
            values[i].extend([i])
        correct = [row[:-1] for row in sorted(values)]
        result = [
            [tr.header.__getattribute__(self.attr1), tr.header.__getattribute__(self.attr2)] for tr in self.sgy.traces
        ]
        self.assertEqual(result, correct)