diff --git a/python/sunbeam/schedule.py b/python/sunbeam/schedule.py index 33100b569..bb48b4028 100644 --- a/python/sunbeam/schedule.py +++ b/python/sunbeam/schedule.py @@ -85,10 +85,10 @@ class Group(object): def __getitem__(self, name): return Group(self._schedule._group(name), self._schedule, self.timestep) - def wells(self, timestep): - names = self._wellnames(timestep) - wells = {well.name: well for well in self._schedule.wells} - return map(wells.__getitem__, filter(wells.__contains__, names)) + @property + def wells(self): + names = self._wellnames(self.timestep) + return [w for w in self._schedule.wells if w.name in names] @property def parent(self): @@ -100,8 +100,8 @@ class Group(object): @property def children(self): - l = [] chl = self._schedule._group_tree(self.timestep)._children(self.name) - for elem in chl: - l.append(Group(self._schedule._group(elem), self._schedule, self.timestep)) - return l + g = lambda elt : Group(self._schedule._group(elt), + self._schedule, + self.timestep) + return [g(elem) for elem in chl] diff --git a/tests/group_tree.py b/tests/group_tree.py index 3e12775be..26c98c8fd 100644 --- a/tests/group_tree.py +++ b/tests/group_tree.py @@ -18,15 +18,18 @@ class TestGroupTree(unittest.TestCase): self.assertEqual(None, gr.parent.parent) def test_timestep_groups(self): + total = 0 for group in self.es.schedule.groups(timestep=3): for child in group.children: self.assertIsNotNone(child.name) + total += 1 + self.assertEqual(13, total) group = self.es.schedule.group(timestep=3)['PROD'] children = ['MANI-B1', 'MANI-B2', 'MANI-D1', 'MANI-D2', 'MANI-E1', 'MANI-E2'] - for child in group.children: - self.assertIn(child.name, children) - + names = [child.name for child in group.children] + self.assertEqual(sunbeam.schedule.Group, type(child)) + self.assertEqual(set(children), set(names)) if __name__ == '__main__': unittest.main() diff --git a/tests/schedule.py b/tests/schedule.py index 81519536d..b59e40bb0 100644 --- a/tests/schedule.py +++ b/tests/schedule.py @@ -34,7 +34,7 @@ class TestSchedule(unittest.TestCase): self.assertEqual(dt.date(2015, 12, 31), timesteps[7]) def testGroups(self): - g1 = self.sch.group()['G1'].wells(0) + g1 = self.sch.group()['G1'].wells self.assertEqual(2, len(g1)) def head(xs): return next(iter(xs))