custom filename in write_to_csv + some codestyle from pycharm

Author: fobiasmog

html_table_extractor/extractor.py

@@ -69,8 +69,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

@@ -2,7 +2,10 @@
 # -*- coding: utf-8 -*-
 
 import unittest
+from unittest.mock import patch, mock_open
+
 from bs4 import BeautifulSoup
+
 from html_table_extractor.extractor import Extractor
 
 
@@ -132,5 +135,34 @@ def test_return_list(self):
         )
 
 
+@patch('csv.writer')
+@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_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()