diff --git a/lib.py b/lib.py
index 6bc7f6d..f482218 100644
--- a/lib.py
+++ b/lib.py
@@ -3,14 +3,18 @@ WORD_LENGTH = 5
 ROUNDS = 6
 INSTRUCTIONS = "Positional letters (green) are CAPITAL, non-positional (yellow) are small, others are blanks e.g.  SHe-r"
 BLANK = "-"
+DICTIONARY = '/usr/share/dict/words'
 
 
-def cat_words():
-    with open('/usr/share/dict/words') as f:
-        for word in f.readlines():
-            # remove end of line
-            word = word[:-1]
-            # filter out five letter words
-            if len(word) == WORD_LENGTH and "'" not in word and not word[0].isupper():
-                yield word
-
+def cat_words(dictionary=DICTIONARY):
+    if type(dictionary) is list:
+        for word in dictionary:
+            yield word
+    else:
+        with open(dictionary) as f:
+            for word in f.readlines():
+                # remove end of line
+                word = word[:-1]
+                # filter out five letter words
+                if len(word) == WORD_LENGTH and "'" not in word and not word[0].isupper():
+                    yield word
diff --git a/play.py b/play.py
index f95cb66..3cd2996 100644
--- a/play.py
+++ b/play.py
@@ -47,10 +47,6 @@ def main_loop():
     print("Word was: {}".format(word))
 
 
-def main():
-    main_loop()
-
-
 def test():
     assert build_letter_hint('yummy', 'slily') == '----Y'
     assert build_letter_hint('yummy', 'slyly') == 'y---Y'
@@ -59,6 +55,10 @@ def test():
     assert build_letter_hint('mmmyy', 'slyyy') == '---YY'
 
 
-if __name__ == "__main__":
+def main():
     test()
+    main_loop()
+
+
+if __name__ == "__main__":
     main()
diff --git a/solve.py b/solve.py
index ed01600..36366fc 100644
--- a/solve.py
+++ b/solve.py
@@ -1,6 +1,6 @@
 import random
 
-from lib import cat_words, INSTRUCTIONS, BLANK, ROUNDS, DEBUG
+from lib import cat_words, INSTRUCTIONS, BLANK, ROUNDS, DEBUG, DICTIONARY
 
 
 def remove_non_letters(word, non_letters):
@@ -32,8 +32,8 @@ def remove_non_positional_letters(word, non_positional_letters):
     return False
 
 
-def round_words(non_letters, all_positional_letters, all_non_positional_letters):
-    for word in cat_words():
+def round_words(non_letters, all_positional_letters, all_non_positional_letters, dictionary=DICTIONARY):
+    for word in cat_words(dictionary):
         skip = remove_non_letters(word, non_letters)
         if skip:
             continue
@@ -135,13 +135,22 @@ def main_loop():
                 break
         if all([letter.isupper for letter in letters]):
             print("WIN")
+            break
     else:
         print("Sorry. Remaining words: {}".format(len(this_round_words)))
         for word in this_round_words:
             print(word)
 
 
+def test():
+    assert [w for w in round_words('', ['----y'], ['-----'], dictionary=['slily'])] == ['slily']
+    assert [w for w in round_words('', ['----y'], ['y----'], dictionary=['slyly'])] == ['slyly']
+    assert [w for w in round_words('', ['---yy'], ['y----'], dictionary=['slyyy'])] == ['slyyy']
+    assert [w for w in round_words('', ['---yy'], ['-----'], dictionary=['slyyy'])] == ['slyyy']
+
+
 def main():
+    test()
     main_loop()