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