diff --git a/tests/test_random_word.py b/tests/test_random_word.py index 6b8d769..29d49bb 100644 --- a/tests/test_random_word.py +++ b/tests/test_random_word.py @@ -151,6 +151,12 @@ def test_random_words_not_enough_with_return_less_if_necessary(self): ] ) + def test_random_words_enough_with_return_less_if_necessary(self): + """Test to see if setting return_less_if_necessary doesn't do anything if there are enough words.""" + data_1 = self.rw.random_words(5, starts_with="mo", return_less_if_necessary=True) + data_2 = self.rw.random_words(5, starts_with="mo", return_less_if_necessary=False) + assert len(data_1) == len(data_2) == 5 + def test_word(self): """Test the word method""" data = self.rw.word() diff --git a/wonderwords/random_word.py b/wonderwords/random_word.py index 133e969..8fcfab0 100644 --- a/wonderwords/random_word.py +++ b/wonderwords/random_word.py @@ -391,14 +391,15 @@ def random_words( exclude_with_spaces=exclude_with_spaces, ) - if not return_less_if_necessary and len(choose_from) < amount: - raise NoWordsToChoseFrom( - "There aren't enough words to choose from. Cannot generate " - f"{str(amount)} word(s)" - ) - elif return_less_if_necessary: - random.shuffle(choose_from) - return choose_from + if len(choose_from) < amount: + if return_less_if_necessary: + random.shuffle(choose_from) + return choose_from + else: + raise NoWordsToChoseFrom( + "There aren't enough words to choose from. Cannot generate " + f"{str(amount)} word(s)" + ) words = [] for _ in range(amount):