diff options
-rw-r--r-- | factory/fuzzy.py | 5 | ||||
-rw-r--r-- | tests/test_fuzzy.py | 22 |
2 files changed, 19 insertions, 8 deletions
diff --git a/factory/fuzzy.py b/factory/fuzzy.py index 34949c5..2ea544a 100644 --- a/factory/fuzzy.py +++ b/factory/fuzzy.py @@ -107,18 +107,19 @@ class FuzzyChoice(BaseFuzzyAttribute): class FuzzyInteger(BaseFuzzyAttribute): """Random integer within a given range.""" - def __init__(self, low, high=None, **kwargs): + def __init__(self, low, high=None, step=1, **kwargs): if high is None: high = low low = 0 self.low = low self.high = high + self.step = step super(FuzzyInteger, self).__init__(**kwargs) def fuzz(self): - return random.randint(self.low, self.high) + return random.randrange(self.low, self.high + 1, self.step) class FuzzyDecimal(BaseFuzzyAttribute): diff --git a/tests/test_fuzzy.py b/tests/test_fuzzy.py index d6f33bb..1caeb0a 100644 --- a/tests/test_fuzzy.py +++ b/tests/test_fuzzy.py @@ -89,24 +89,34 @@ class FuzzyIntegerTestCase(unittest.TestCase): self.assertIn(res, [0, 1, 2, 3, 4]) def test_biased(self): - fake_randint = lambda low, high: low + high + fake_randrange = lambda low, high, step: (low + high) * step fuzz = fuzzy.FuzzyInteger(2, 8) - with mock.patch('random.randint', fake_randint): + with mock.patch('random.randrange', fake_randrange): res = fuzz.evaluate(2, None, False) - self.assertEqual(10, res) + self.assertEqual((2 + 8 + 1) * 1, res) def test_biased_high_only(self): - fake_randint = lambda low, high: low + high + fake_randrange = lambda low, high, step: (low + high) * step fuzz = fuzzy.FuzzyInteger(8) - with mock.patch('random.randint', fake_randint): + with mock.patch('random.randrange', fake_randrange): + res = fuzz.evaluate(2, None, False) + + self.assertEqual((0 + 8 + 1) * 1, res) + + def test_biased_with_step(self): + fake_randrange = lambda low, high, step: (low + high) * step + + fuzz = fuzzy.FuzzyInteger(5, 8, 3) + + with mock.patch('random.randrange', fake_randrange): res = fuzz.evaluate(2, None, False) - self.assertEqual(8, res) + self.assertEqual((5 + 8 + 1) * 3, res) class FuzzyDecimalTestCase(unittest.TestCase): |