Mock stdin wordlist

Author: gigkokman

lib/helpers/wordlist_helper.py

@@ -19,7 +19,6 @@ def get_wordlist(self, wordlist_files=None):
         wordlist = []
         wordlist_types = []
         stdin_words = self.get_stdin_wordlist()
-        print(stdin_words)
         if stdin_words:
             wordlist_types.append('stdin')
             wordlist.extend(stdin_words)

tests/helpers/test_wordlist_helper.py

@@ -1,5 +1,4 @@
 import unittest
-from mock import Mock
 from mock import patch
 
 from lib.helpers.wordlist_helper import WordList
@@ -13,16 +12,16 @@ def setUp(self):
             self.default_wordlist = list(word_file.read().splitlines())
 
     def test_get_wordlist_from_stdin(self):
-        stdin_list = ['keyword1', 'keyword1']
+        stdin_wordlist = ['keyword1', 'keyword1']
         expected_wordlist = []
-        expected_wordlist.extend(stdin_list)
+        expected_wordlist.extend(stdin_wordlist)
         expected_wordlist.extend(self.default_wordlist)
-        with patch('sys.stdin') as mock_stdin:
-            mock_stdin.read = Mock(return_value='\n'.join(stdin_list))
-            mock_stdin.isatty = Mock(return_value=False)
+        with patch('lib.helpers.wordlist_helper.WordList.get_stdin_wordlist', return_value=stdin_wordlist):
             wordlist, wordlist_types = self.wordlist.get_wordlist()
             self.assertEqual(wordlist, expected_wordlist)
 
     def test_using_default_wordlist(self):
-        wordlist, wordlist_types = self.wordlist.get_wordlist()
-        self.assertEqual(wordlist, self.default_wordlist)
+        stdin_wordlist = []
+        with patch('lib.helpers.wordlist_helper.WordList.get_stdin_wordlist', return_value=stdin_wordlist):
+            wordlist, wordlist_types = self.wordlist.get_wordlist()
+            self.assertEqual(wordlist, self.default_wordlist)