244 lines
12 KiB
Python
244 lines
12 KiB
Python
import logging
|
|
from pathlib import Path
|
|
from unittest import IsolatedAsyncioTestCase
|
|
from unittest.mock import patch, MagicMock, AsyncMock, call
|
|
|
|
from bs4 import BeautifulSoup
|
|
|
|
from mwfin import functions
|
|
from mwfin.constants import HTML_PARSER, BASE_URL, FIN_STMT_URL_SUFFIX, IS, BS, CF, END_DATE
|
|
|
|
|
|
THIS_DIR = Path(__file__).parent
|
|
|
|
|
|
class FunctionsTestCase(IsolatedAsyncioTestCase):
|
|
# boiled down & accurate structure of a relevant data table
|
|
# https://www.marketwatch.com/investing/stock/aapl/financials/cash-flow
|
|
# view page source @ line 2055
|
|
TEST_HTML_FILE_PATH = Path(THIS_DIR, 'test_structure.html')
|
|
|
|
log_lvl: int
|
|
test_html: str
|
|
|
|
@staticmethod
|
|
def get_mock_session(response_text: str = None) -> MagicMock:
|
|
mock_response = MagicMock()
|
|
mock_response.text = AsyncMock(return_value=response_text)
|
|
mock_get_return = MagicMock()
|
|
mock_get_return.__aenter__ = AsyncMock(return_value=mock_response)
|
|
mock_session_obj = MagicMock()
|
|
mock_session_obj.get = MagicMock(return_value=mock_get_return)
|
|
return mock_session_obj
|
|
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
with open(cls.TEST_HTML_FILE_PATH, 'r') as f:
|
|
cls.test_html = f.read()
|
|
cls.test_soup = BeautifulSoup(cls.test_html, HTML_PARSER)
|
|
cls.log_lvl = functions.log.level
|
|
functions.log.setLevel(logging.CRITICAL)
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
functions.log.setLevel(cls.log_lvl)
|
|
|
|
@patch.object(functions, 'ClientSession')
|
|
async def test_soup_from_url(self, mock_session_cls):
|
|
test_html = '<b>foo</b>'
|
|
mock_session_cls.return_value = mock_session_obj = self.get_mock_session(test_html)
|
|
expected_output = BeautifulSoup(test_html, 'html.parser')
|
|
output = await functions.soup_from_url('baz', mock_session_obj)
|
|
self.assertEqual(expected_output, output)
|
|
|
|
def test_get_row_indent(self):
|
|
mock_row = BeautifulSoup('<tr><div>foo</div></tr>', HTML_PARSER).tr
|
|
expected_output = 0
|
|
output = functions.get_row_indent(mock_row)
|
|
self.assertEqual(expected_output, output)
|
|
trs = self.test_soup.find_all('tr')
|
|
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)
|
|
|
|
@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')
|
|
@patch.object(functions, 'extract_end_dates')
|
|
def test_extract_all_data(self, mock_extract_end_dates, mock_get_all_table_rows, mock_extract_row_data):
|
|
test_end_dates = ('foo', 'bar')
|
|
mock_extract_end_dates.return_value = test_end_dates
|
|
test_relevant_rows = ['tr1', 'tr2']
|
|
mock_get_all_table_rows.return_value = test_relevant_rows
|
|
test_row_data = ('item_name', (123, 456))
|
|
mock_extract_row_data.return_value = test_row_data
|
|
expected_output = {
|
|
END_DATE: test_end_dates,
|
|
test_row_data[0]: test_row_data[1],
|
|
test_row_data[0]: test_row_data[1],
|
|
}
|
|
output = functions.extract_all_data(self.test_soup)
|
|
self.assertDictEqual(expected_output, output)
|
|
mock_extract_end_dates.assert_called_once_with(self.test_soup)
|
|
mock_get_all_table_rows.assert_called_once_with(self.test_soup)
|
|
mock_extract_row_data.assert_has_calls([call(test_relevant_rows[0]), call(test_relevant_rows[1])])
|
|
|
|
@patch.object(functions, 'extract_all_data')
|
|
@patch.object(functions, 'soup_from_url')
|
|
async def test__get_single_company_fin_stmt(self, mock_soup_from_url, mock_extract_all_data):
|
|
mock_session = MagicMock()
|
|
test_ticker, statement = 'bar', BS
|
|
test_url = f'{BASE_URL}/{test_ticker}/financials{FIN_STMT_URL_SUFFIX[statement]}'
|
|
mock_soup_from_url.return_value = mock_soup = MagicMock()
|
|
mock_extract_all_data.return_value = expected_output = {'foo': 'bar'}
|
|
|
|
quarterly = False
|
|
output = await functions._get_single_company_fin_stmt(statement, test_ticker, quarterly, mock_session)
|
|
self.assertDictEqual(expected_output, output)
|
|
mock_soup_from_url.assert_called_once_with(test_url, mock_session)
|
|
mock_extract_all_data.assert_called_once_with(mock_soup)
|
|
mock_soup_from_url.reset_mock()
|
|
mock_extract_all_data.reset_mock()
|
|
|
|
quarterly = True
|
|
output = await functions._get_single_company_fin_stmt(statement, test_ticker, quarterly, mock_session)
|
|
self.assertDictEqual(expected_output, output)
|
|
mock_soup_from_url.assert_called_once_with(test_url + '/quarter', mock_session)
|
|
mock_extract_all_data.assert_called_once_with(mock_soup)
|
|
|
|
@patch.object(functions, '_get_single_company_fin_stmt')
|
|
async def test__get_multi_companies_fin_stmt(self, mock__get_single_company_fin_stmt):
|
|
statement, sym1, sym2, quarterly, mock_session = 'xyz', 'foo', 'bar', False, MagicMock()
|
|
mock__get_single_company_fin_stmt.return_value = expected_output = 'baz'
|
|
output = await functions._get_multi_companies_fin_stmt(statement, sym1,
|
|
quarterly=quarterly, session=mock_session)
|
|
self.assertEqual(expected_output, output)
|
|
mock__get_single_company_fin_stmt.assert_called_once_with(statement, sym1, quarterly, mock_session)
|
|
mock__get_single_company_fin_stmt.reset_mock()
|
|
|
|
expected_output = {sym1: expected_output, sym2: expected_output}
|
|
output = await functions._get_multi_companies_fin_stmt(statement, sym1, sym2,
|
|
quarterly=quarterly, session=mock_session)
|
|
self.assertDictEqual(expected_output, output)
|
|
mock__get_single_company_fin_stmt.assert_has_calls([
|
|
call(statement, sym1, quarterly, mock_session),
|
|
call(statement, sym2, quarterly, mock_session)
|
|
])
|
|
|
|
async def _helper_test_get_any_statement(self, stmt: str, mock__get_multi_companies_fin_stmt):
|
|
sym1, sym2, quarterly, batch_size, mock_session = 'foo', 'bar', False, 2, MagicMock()
|
|
mock__get_multi_companies_fin_stmt.return_value = expected_output = 'baz'
|
|
if stmt == BS:
|
|
function = functions.get_balance_sheet
|
|
elif stmt == IS:
|
|
function = functions.get_income_statement
|
|
elif stmt == CF:
|
|
function = functions.get_cash_flow_statement
|
|
else:
|
|
raise ValueError
|
|
output = await function(sym1, sym2, quarterly=quarterly, concurrent_batch_size=batch_size, session=mock_session)
|
|
self.assertEqual(expected_output, output)
|
|
mock__get_multi_companies_fin_stmt.assert_called_once_with(
|
|
stmt, sym1, sym2, quarterly=quarterly, concurrent_batch_size=batch_size, session=mock_session
|
|
)
|
|
|
|
@patch.object(functions, '_get_multi_companies_fin_stmt')
|
|
async def test_get_balance_sheet(self, mock__get_multi_companies_fin_stmt):
|
|
await self._helper_test_get_any_statement(BS, mock__get_multi_companies_fin_stmt)
|
|
|
|
@patch.object(functions, '_get_multi_companies_fin_stmt')
|
|
async def test_get_income_statement(self, mock__get_multi_companies_fin_stmt):
|
|
await self._helper_test_get_any_statement(IS, mock__get_multi_companies_fin_stmt)
|
|
|
|
@patch.object(functions, '_get_multi_companies_fin_stmt')
|
|
async def test_get_cash_flow_statement(self, mock__get_multi_companies_fin_stmt):
|
|
await self._helper_test_get_any_statement(CF, mock__get_multi_companies_fin_stmt)
|
|
|
|
@patch.object(functions, '_get_single_company_fin_stmt')
|
|
async def test__get_single_company_all_financials(self, mock__get_single_company_fin_stmt):
|
|
symbol, quarterly, mock_session = 'foo', False, MagicMock()
|
|
mock__get_single_company_fin_stmt.return_value = bar = 'bar'
|
|
expected_output = {BS: bar, IS: bar, CF: bar}
|
|
output = await functions._get_single_company_all_financials(symbol, quarterly, mock_session)
|
|
self.assertDictEqual(expected_output, output)
|
|
mock__get_single_company_fin_stmt.assert_has_calls([
|
|
call(BS, symbol, quarterly, mock_session),
|
|
call(IS, symbol, quarterly, mock_session),
|
|
call(CF, symbol, quarterly, mock_session)
|
|
])
|
|
|
|
@patch.object(functions, '_get_single_company_all_financials')
|
|
async def test_get_company_financials(self, mock__get_single_company_all_financials):
|
|
mock__get_single_company_all_financials.return_value = expected_output = 'baz'
|
|
symbol, quarterly, mock_session = 'foo', False, MagicMock()
|
|
output = await functions.get_all_financials(symbol, quarterly=quarterly, session=mock_session)
|
|
self.assertEqual(expected_output, output)
|
|
mock__get_single_company_all_financials.assert_called_once_with(symbol, quarterly, mock_session)
|
|
mock__get_single_company_all_financials.reset_mock()
|
|
|
|
test_sym1, test_sym2 = 'x', 'y'
|
|
expected_output = {test_sym1: expected_output, test_sym2: expected_output}
|
|
output = await functions.get_all_financials(test_sym1, test_sym2,
|
|
quarterly=quarterly, session=mock_session)
|
|
self.assertDictEqual(expected_output, output)
|
|
mock__get_single_company_all_financials.assert_has_calls([
|
|
call(test_sym1, quarterly, mock_session),
|
|
call(test_sym2, quarterly, mock_session)
|
|
])
|
|
|
|
@patch.object(functions, 'ClientSession')
|
|
async def test_integration_get_company_financials(self, mock_session_cls):
|
|
mock_session_cls.return_value = mock_session_obj = self.get_mock_session(self.test_html)
|
|
symbol = 'foo'
|
|
# Since the web request is mocked we always receive the same HTML markup.
|
|
expected_output = {
|
|
BS: {
|
|
END_DATE: (0, 'End_Date_1', 'End_Date_2'),
|
|
'foo': (1, 1., -2.),
|
|
'bar': (2, 2., -3.),
|
|
'baz': (3, 3., -4.)
|
|
},
|
|
IS: {
|
|
END_DATE: (0, 'End_Date_1', 'End_Date_2'),
|
|
'foo': (1, 1., -2.),
|
|
'bar': (2, 2., -3.),
|
|
'baz': (3, 3., -4.)
|
|
},
|
|
CF: {
|
|
END_DATE: (0, 'End_Date_1', 'End_Date_2'),
|
|
'foo': (1, 1., -2.),
|
|
'bar': (2, 2., -3.),
|
|
'baz': (3, 3., -4.)
|
|
}
|
|
}
|
|
output = await functions.get_all_financials(symbol, session=mock_session_obj)
|
|
self.assertDictEqual(expected_output, output)
|
|
mock_session_obj.get.assert_has_calls([
|
|
call(f'{BASE_URL}/{symbol}/financials{FIN_STMT_URL_SUFFIX[BS]}'),
|
|
call(f'{BASE_URL}/{symbol}/financials{FIN_STMT_URL_SUFFIX[IS]}'),
|
|
call(f'{BASE_URL}/{symbol}/financials{FIN_STMT_URL_SUFFIX[CF]}'),
|
|
])
|