Test providing STDIN and wordlist

Author: gigkokman

lib/helpers/wordlist_helper.py

@@ -20,18 +20,19 @@ def get_stdin_wordlist(self):
             if not sys.stdin.isatty() else []
 
     def get_wordlist(self, wordlist_files=None):
-        stdin_words = self.get_stdin_wordlist()
         default_wordlist_file = DEFAULT_WORDLIST_FILE
+
+        stdin_words = self.get_stdin_wordlist()
         if stdin_words:
-            self.set_words('stdin', stdin_words)
+            self.set_words(words_type='stdin', words=stdin_words)
             default_wordlist_file = None
 
         combined_files = wordlist_files or default_wordlist_file
         combined = get_combined_word_lists(combined_files)
         if combined:
             words_type = 'wordlists: {}'.format(
                 ', '.join(combined['file_paths']))
-            self.set_words(words_type, combined['words'])
+            self.set_words(words_type=words_type, words=combined['words'])
 
         return self.wordlist, self.wordlist_types
 

tests/helpers/test_wordlist_helper.py

@@ -1,10 +1,22 @@
 import unittest
+import pytest
 from mock import patch
 
 from lib.helpers.wordlist_helper import WordList
 from lib.helpers.wordlist_helper import DEFAULT_WORDLIST_FILE
 
 
+@pytest.fixture(scope='class')
+def user_wordlist(request, tmpdir_factory):
+    user_wordlist = ['user-word1', 'user-word2']
+    tmpdir = tmpdir_factory.mktemp('user_wordlist')
+    user_wordlist_file = tmpdir.join('user-wordlist.txt')
+    user_wordlist_file.write('\n'.join(user_wordlist))
+    request.cls.user_wordlist_file = str(user_wordlist_file)
+    request.cls.user_wordlist = user_wordlist
+
+
+@pytest.mark.usefixtures('user_wordlist')
 class TestWordList(unittest.TestCase):
     def setUp(self):
         self.wordlist = WordList()
@@ -19,6 +31,15 @@ def test_get_wordlist_from_stdin(self):
             wordlist, wordlist_types = self.wordlist.get_wordlist()
             self.assertEqual(wordlist, expected_wordlist)
 
+    def test_get_wordlist_from_stdin_and_wordlist(self):
+        stdin_wordlist = ['keyword1', 'keyword1']
+        expected_wordlist = []
+        expected_wordlist.extend(stdin_wordlist)
+        expected_wordlist.extend(self.user_wordlist)
+        with patch('lib.helpers.wordlist_helper.WordList.get_stdin_wordlist', return_value=stdin_wordlist):
+            wordlist, wordlist_types = self.wordlist.get_wordlist(self.user_wordlist_file)
+            self.assertEqual(wordlist, expected_wordlist)
+
     def test_using_default_wordlist(self):
         stdin_wordlist = []
         with patch('lib.helpers.wordlist_helper.WordList.get_stdin_wordlist', return_value=stdin_wordlist):