diff --git a/tests/test_functions.py b/tests/test_functions.py index 4c503e5..2cea0f8 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -66,8 +66,9 @@ class FunctionsTestCase(IsolatedAsyncioTestCase): with self.assertRaises(UnknownFinancialStatementItem): functions.is_relevant_table_row(test_soup.tr) + @patch.object(functions, 'open') @patch.object(functions, 'is_relevant_table_row') - def test_find_relevant_table_rows(self, mock_is_relevant_table_row): + def test_find_relevant_table_rows(self, mock_is_relevant_table_row, mock_open): mock_is_relevant_table_row.return_value = True expected_output = self.test_soup.find('div', attrs={'class': 'financials'}).tbody.find_all('tr') tr0, tr1 = expected_output @@ -81,6 +82,14 @@ class FunctionsTestCase(IsolatedAsyncioTestCase): output = functions.find_relevant_table_rows(self.test_soup) self.assertListEqual(expected_output, output) mock_is_relevant_table_row.assert_has_calls([call(tr0), call(tr1)]) + mock_is_relevant_table_row.reset_mock() + + mock_write = mock_open.return_value.__enter__.return_value.write + with patch.object(functions, 'DEV_MODE', new=True): + output = functions.find_relevant_table_rows(self.test_soup) + self.assertListEqual(expected_output, output) + mock_is_relevant_table_row.assert_has_calls([call(tr0), call(tr1)]) + mock_write.assert_has_calls([call(str(self.test_soup)), call(str(self.test_soup))]) def test_extract_row_data(self): test_row = self.test_soup.find('div', attrs={'class': 'financials'}).tbody.tr