diff --git a/rolltable/tables.py b/rolltable/tables.py index acab4dc..d8f16d2 100644 --- a/rolltable/tables.py +++ b/rolltable/tables.py @@ -128,10 +128,7 @@ class DataSource: ] else: # If the option is either a list or a string, just select it. - if rand: - choices = [random.choice(self.data[option])] - else: - choices = self.data[option] + choices = self.data[option] for choice in choices: # If the randomly-selected choice is a dict, choose a random item and return a list consisting @@ -153,6 +150,7 @@ class DataSource: flattened.append(choice) return flattened + class RollTable: """ Generate a roll table using weighted distributions of random options. diff --git a/tests/test_tables.py b/tests/test_tables.py index 43dd6f7..6e1ccb9 100644 --- a/tests/test_tables.py +++ b/tests/test_tables.py @@ -204,7 +204,14 @@ def test_text(): assert repr(tables.RollTable([fixture_source])) -def test_as_dict(): - source = tables.RollTable([fixture_no_descriptions]).datasources[0] - ds = source.as_dict() - assert ds['option 1']['choice'] == 'choice 1' +@pytest.mark.parametrize('table', [ + tables.RollTable([fixture_no_options]), + tables.RollTable([fixture_one_choice]), + tables.RollTable([fixture_metadata + fixture_source]), + tables.RollTable([fixture_source]), + tables.RollTable([fixture_no_options]), + tables.RollTable([fixture_lists_and_dicts]), +]) +def test_as_dict(table): + for src in table.datasources: + assert src.as_dict()