mwfin/tests/test_functions.py

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_financial_statement(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_financial_statement(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_financial_statement(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)

    async def _helper_test_get_any_statement(self, statement: str, mock__get_financial_statement):
        symbol, quarterly, mock_session = 'foo', False, MagicMock()
        mock__get_financial_statement.return_value = expected_output = 'bar'
        if statement == BS:
            function = functions.get_balance_sheet
        elif statement == IS:
            function = functions.get_income_statement
        elif statement == CF:
            function = functions.get_cash_flow_statement
        else:
            raise ValueError
        output = await function(symbol, quarterly=quarterly, session=mock_session)
        self.assertEqual(expected_output, output)
        mock__get_financial_statement.assert_called_once_with(statement, symbol, quarterly, mock_session)
        mock__get_financial_statement.reset_mock()
        
        symbol1, symbol2 = 'x', 'y'
        expected_output = {symbol1: expected_output, symbol2: expected_output}
        output = await function(symbol1, symbol2, quarterly=quarterly, session=mock_session)
        self.assertDictEqual(expected_output, output)
        mock__get_financial_statement.assert_has_calls([
            call(statement, symbol1, quarterly, mock_session),
            call(statement, symbol2, quarterly, mock_session),
        ])

    @patch.object(functions, '_get_financial_statement')
    async def test_get_balance_sheet(self, mock__get_financial_statement):
        await self._helper_test_get_any_statement(BS, mock__get_financial_statement)

    @patch.object(functions, '_get_financial_statement')
    async def test_get_income_statement(self, mock__get_financial_statement):
        await self._helper_test_get_any_statement(IS, mock__get_financial_statement)

    @patch.object(functions, '_get_financial_statement')
    async def test_get_cash_flow_statement(self, mock__get_financial_statement):
        await self._helper_test_get_any_statement(CF, mock__get_financial_statement)

    @patch.object(functions, 'get_cash_flow_statement')
    @patch.object(functions, 'get_income_statement')
    @patch.object(functions, 'get_balance_sheet')
    async def test_get_company_financials(self, mock_get_bs, mock_get_is, mock_get_cf):
        mock_end_dates = ('bar', 'baz')
        mock_get_bs.return_value = {END_DATE: mock_end_dates, 'a': (1, 2)}
        mock_get_is.return_value = {END_DATE: mock_end_dates, 'b': (2, 3)}
        mock_get_cf.return_value = {END_DATE: mock_end_dates, 'c': (3, 4)}
        expected_output = {
            BS: {END_DATE: mock_end_dates, 'a': (1, 2)},
            IS: {END_DATE: mock_end_dates, 'b': (2, 3)},
            CF: {END_DATE: mock_end_dates, 'c': (3, 4)}
        }
        symbol, quarterly, mock_session = 'foo', False, MagicMock()
        output = await functions.get_company_financials(symbol, quarterly=quarterly, session=mock_session)
        self.assertDictEqual(expected_output, output)
        mock_get_bs.assert_called_once_with(symbol, quarterly=quarterly, session=mock_session)
        mock_get_is.assert_called_once_with(symbol, quarterly=quarterly, session=mock_session)
        mock_get_cf.assert_called_once_with(symbol, quarterly=quarterly, session=mock_session)
        mock_get_bs.reset_mock()
        mock_get_is.reset_mock()
        mock_get_cf.reset_mock()

        test_symbol1, test_symbol2 = 'x', 'y'
        expected_output = {test_symbol1: expected_output, test_symbol2: expected_output}
        output = await functions.get_company_financials(test_symbol1, test_symbol2,
                                                        quarterly=quarterly, session=mock_session)
        self.assertDictEqual(expected_output, output)
        mock_get_bs.assert_has_calls([
            call(test_symbol1, quarterly=quarterly, session=mock_session),
            call(test_symbol2, quarterly=quarterly, session=mock_session)
        ])
        mock_get_is.assert_has_calls([
            call(test_symbol1, quarterly=quarterly, session=mock_session),
            call(test_symbol2, quarterly=quarterly, session=mock_session)
        ])
        mock_get_cf.assert_has_calls([
            call(test_symbol1, quarterly=quarterly, session=mock_session),
            call(test_symbol2, quarterly=quarterly, session=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_company_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]}'),
        ])