diff --git a/tests/test_functions.py b/tests/test_functions.py index 4af897b..588e82d 100644 --- a/tests/test_functions.py +++ b/tests/test_functions.py @@ -51,21 +51,36 @@ class FunctionsTestCase(IsolatedAsyncioTestCase): output = await functions.soup_from_url('baz', mock_session_obj) self.assertEqual(expected_output, output) - def test_extract_end_dates(self): - expected_output = ('End_Date_1', 'End_Date_2') + def test_get_row_indent(self): + trs = self.test_soup.find_all('tr') + expected_output = 0 + output = functions.get_row_indent(trs[0]) + self.assertEqual(expected_output, output) + for i, tr in enumerate(trs[1:], start=1): + output = functions.get_row_indent(tr) + self.assertEqual(i, output) + + @patch.object(functions, 'get_row_indent') + def test_extract_end_dates(self, mock_get_row_indent): + mock_get_row_indent.return_value = 0 + expected_output = (0, 'End_Date_1', 'End_Date_2') output = functions.extract_end_dates(self.test_soup) self.assertTupleEqual(expected_output, output) + mock_get_row_indent.assert_called_once_with(self.test_soup.tr) def test_get_all_table_rows(self): expected_output = self.test_soup.find('div', attrs={'class': 'financials'}).tbody.find_all('tr') output = functions.get_all_table_rows(self.test_soup) self.assertSequenceEqual(expected_output, output) - def test_extract_row_data(self): + @patch.object(functions, 'get_row_indent') + def test_extract_row_data(self, mock_get_row_indent): + mock_get_row_indent.return_value = 1 test_row = self.test_soup.find('div', attrs={'class': 'financials'}).tbody.tr expected_output = ('foo', (1, 1., -2.)) output = functions.extract_row_data(test_row) self.assertTupleEqual(expected_output, output) + mock_get_row_indent.assert_called_once_with(test_row) @patch.object(functions, 'extract_row_data') @patch.object(functions, 'get_all_table_rows')