mwfin/tests/test_functions.py

222 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_extract_end_dates(self):
expected_output = ('End_Date_1', 'End_Date_2')
output = functions.extract_end_dates(self.test_soup)
self.assertTupleEqual(expected_output, output)
def test_is_relevant_table_row(self):
test_soup = BeautifulSoup('<tr><td><div> Cash & Short Term Investments </div></td></tr>', HTML_PARSER)
self.assertTrue(functions.is_relevant_table_row(test_soup.tr))
test_soup = BeautifulSoup('<tr><td><div> Cash & Short Term Investments Growth </div></td></tr>', HTML_PARSER)
self.assertFalse(functions.is_relevant_table_row(test_soup.tr))
test_soup = BeautifulSoup('<tr><td><div> baz </div></td></tr>', HTML_PARSER)
self.assertFalse(functions.is_relevant_table_row(test_soup.tr))
@patch.object(functions, 'is_relevant_table_row')
def test_find_relevant_table_rows(self, mock_is_relevant_table_row):
mock_is_relevant_table_row.return_value = True
expected_output = self.test_soup.find('div', attrs={'class': 'financials'}).tbody.find_all('tr')
output = functions.find_relevant_table_rows(self.test_soup)
self.assertListEqual(expected_output, output)
mock_is_relevant_table_row.assert_has_calls([call(expected_output[0]), call(expected_output[1])])
def test_extract_row_data(self):
test_row = self.test_soup.find('div', attrs={'class': 'financials'}).tbody.tr
expected_output = ('Cash & Short Term Investments', (11000000, -22000000))
output = functions.extract_row_data(test_row)
self.assertTupleEqual(expected_output, output)
@patch.object(functions, 'extract_row_data')
@patch.object(functions, 'find_relevant_table_rows')
@patch.object(functions, 'extract_end_dates')
def test_extract_all_data(self, mock_extract_end_dates, mock_find_relevant_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_find_relevant_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_find_relevant_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: ('End_Date_1', 'End_Date_2'), 'Cash & Short Term Investments': (11000000, -22000000)},
IS: {END_DATE: ('End_Date_1', 'End_Date_2'), 'Cash & Short Term Investments': (11000000, -22000000)},
CF: {END_DATE: ('End_Date_1', 'End_Date_2'), 'Cash & Short Term Investments': (11000000, -22000000)}
}
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]}'),
])