Search code examples
pythontestingpython-hypothesis

Hypothesis strategy: for each "bucket", draw one value from the bucket


I have the following method for generating random data in one of my tests:

import random

data_categories = {
    'a': [1, 2, 3],
    'b': [4, 5],
    'c': [6, 7, 8]
}

def make_record():
    return [random.choice(vals) for vals in data_categories.values()]

How can I convert this to a Hypothesis strategy?

This is my attempt using hypothesis.strategies.composite, but it's hard to know if I'm using it correctly:

import hypothesis.strategies as hs

@hs.composite
def make_record(draw):
    return [draw(hs.sampled_from(vals)) for vals in data_categories.values()]

Solution

  • Your attempt is basically correct, except that dictionary iteration order may not be reliable - and thus examples may not reproduce or shrink correctly. I'd write it inline, as:

    my_strategy = hs.tuples(
        *[hs.sampled_from(data_categories[k]) for k in sorted(data_categories)]
    ).map(list)