Merge pull request #10 from fobiasmog/master

Custom filename for write_to_csv

Author: yuanxu-li

html_table_extractor/extractor.py

@@ -71,8 +71,8 @@ def parse(self):
     def return_list(self):
         return self._output
 
-    def write_to_csv(self, path='.'):
-        with open(os.path.join(path, 'output.csv'), 'w') as csv_file:
+    def write_to_csv(self, path='.', filename='output.csv'):
+        with open(os.path.join(path, filename), 'w') as csv_file:
             table_writer = csv.writer(csv_file)
             for row in self._output:
                 table_writer.writerow(row)

tests/tests.py

@@ -1,8 +1,16 @@
 #!/usr/bin/python
 # -*- coding: utf-8 -*-
 
+import sys
 import unittest
+
+if sys.version_info[0] < 3:
+    import mock
+else:
+    from unittest import mock
+
 from bs4 import BeautifulSoup
+
 from html_table_extractor.extractor import Extractor
 
 
@@ -132,5 +140,34 @@ def test_return_list(self):
         )
 
 
+@mock.patch('csv.writer')
+@mock.patch('html_table_extractor.extractor.open')
+class TestWriteToCsv(unittest.TestCase):
+    def setUp(self):
+        html = """
+        <table>
+            <tr>
+              <td>1</td>
+              <td>2</td>
+            </tr>
+            <tr>
+              <td>3</td>
+              <td>4</td>
+            </tr>
+        </table>
+        """
+        self.extractor = Extractor(html)
+        self.extractor.parse()
+        mock.mock_open()
+
+    def test_write_to_csv_default(self, csv_mock, _):
+        self.extractor.write_to_csv()
+        csv_mock.assert_called_with('./output.csv', 'w')
+
+    def test_write_to_csv_custom_path_and_filename(self, csv_mock, _):
+        self.extractor.write_to_csv(path='/test/path', filename='testfile.csv')
+        csv_mock.assert_called_with('/test/path/testfile.csv', 'w')
+
+
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()