diff --git a/tools/releasetools/blockimgdiff.py b/tools/releasetools/blockimgdiff.py index 8f06b95f8..69750b28c 100644 --- a/tools/releasetools/blockimgdiff.py +++ b/tools/releasetools/blockimgdiff.py @@ -237,15 +237,23 @@ class Transfer(object): class HeapItem(object): def __init__(self, item): self.item = item - # Negate the score since python's heap is a min-heap and we want - # the maximum score. + # Negate the score since python's heap is a min-heap and we want the + # maximum score. self.score = -item.score + def clear(self): self.item = None + def __bool__(self): - return self.item is None + return self.item is not None + + # Python 2 uses __nonzero__, while Python 3 uses __bool__. + __nonzero__ = __bool__ + + # The rest operations are generated by functools.total_ordering decorator. def __eq__(self, other): return self.score == other.score + def __le__(self, other): return self.score <= other.score diff --git a/tools/releasetools/test_blockimgdiff.py b/tools/releasetools/test_blockimgdiff.py index e5a369478..7084e2153 100644 --- a/tools/releasetools/test_blockimgdiff.py +++ b/tools/releasetools/test_blockimgdiff.py @@ -16,12 +16,43 @@ from __future__ import print_function -import common import unittest -from blockimgdiff import BlockImageDiff, EmptyImage, Transfer +import common +from blockimgdiff import BlockImageDiff, EmptyImage, HeapItem, Transfer from rangelib import RangeSet + +class HealpItemTest(unittest.TestCase): + + class Item(object): + def __init__(self, score): + self.score = score + + def test_init(self): + item1 = HeapItem(self.Item(15)) + item2 = HeapItem(self.Item(20)) + item3 = HeapItem(self.Item(15)) + self.assertTrue(item1) + self.assertTrue(item2) + self.assertTrue(item3) + + self.assertNotEqual(item1, item2) + self.assertEqual(item1, item3) + # HeapItem uses negated scores. + self.assertGreater(item1, item2) + self.assertLessEqual(item1, item3) + self.assertTrue(item1 <= item3) + self.assertFalse(item2 >= item1) + + def test_clear(self): + item = HeapItem(self.Item(15)) + self.assertTrue(item) + + item.clear() + self.assertFalse(item) + + class BlockImageDiffTest(unittest.TestCase): def test_GenerateDigraphOrder(self):