Untitled
unknown
plain_text
6 months ago
4.8 kB
14
Indexable
import unittest from unittest.mock import patch, MagicMock, mock_open import datetime from io import StringIO from blkcore.email_script import get_funds_for_cusip, get_cusips_data_and_send_email, send_email, main class TestEmailScript(unittest.TestCase): @patch('blkdbi.dataobject.DataObject.do_sql') def test_get_funds_for_cusip(self, mock_do_sql): # Setup mock return value for SQL query mock_do_sql.return_value = [MagicMock(fund='fund1'), MagicMock(fund='fund2')] # Define test parameters cusip = "123456789" date = datetime.datetime(2024, 11, 11) # Call the function funds = get_funds_for_cusip(cusip, date) # Validate the behavior mock_do_sql.assert_called_once_with("SELECT * FROM brs_pos WHERE pos_date='11 Nov 2024' AND cusip='123456789'") self.assertEqual(funds, ['fund1', 'fund2']) @patch('blkdbi.dataobject.DataObject.do_sql') @patch('blkcore.email_script.send_email') @patch('builtins.open', new_callable=mock_open, read_data="123456789,987654321") def test_get_cusips_data_and_send_email(self, mock_open, mock_send_email, mock_do_sql): # Setup mock return values mock_do_sql.side_effect = [ [MagicMock(fund='fund1')], # First call (get_funds_for_cusip) [("portfolio1", "full_name1")], # Second call (portfolio info for fund) [("portfolio2", "full_name2")] # Second call (portfolio info for another fund) ] filePath = 'dummy_path.txt' date = datetime.datetime(2024, 11, 11) receiver_emails = ["test@example.com"] # Call the function get_cusips_data_and_send_email(filePath, date, receiver_emails) # Validate that file was read correctly (mock_open) mock_open.assert_called_once_with(filePath) # Check the SQL calls for fetching data mock_do_sql.assert_any_call("SELECT * FROM brs_pos WHERE pos_date='11 Nov 2024' AND cusip='123456789'") mock_do_sql.assert_any_call("SELECT * FROM brs_pos WHERE pos_date='11 Nov 2024' AND cusip='987654321'") mock_do_sql.assert_any_call("SELECT portfolio_name, full_name from portfolios where portfolio_code='fund1'") mock_do_sql.assert_any_call("SELECT portfolio_name, full_name from portfolios where portfolio_code='fund2'") # Verify send_email function is called with the correct arguments mock_send_email.assert_called_once_with( receiver_emails, [ {"Cusip": "123456789", "Portfolio Name": "portfolio1", "Full Name": "full_name1"}, {"Cusip": "987654321", "Portfolio Name": "portfolio2", "Full Name": "full_name2"} ] ) @patch('blkcore.email_script.open', new_callable=mock_open) @patch('blkcore.email_script.BlkSMTP') @patch('blkcore.email_script.get_user', return_value="testuser") def test_send_email(self, mock_get_user, MockBlkSMTP, mock_open): # Prepare the mock for BlkSMTP and open() function mock_smtp_instance = MagicMock() MockBlkSMTP.return_value.__enter__.return_value = mock_smtp_instance receiver_emails = ["test@example.com"] data = [ {"Cusip": "123456789", "Portfolio Name": "portfolio1", "Full Name": "full_name1"}, {"Cusip": "987654321", "Portfolio Name": "portfolio2", "Full Name": "full_name2"} ] # Call the function send_email(receiver_emails, data) # Verify template rendering (check file read) mock_open.assert_called_with('template/template.html') # Verify the email was sent mock_smtp_instance.sendmail.assert_called_once_with( 'testuser@blackrock.com', receiver_emails, mock_smtp_instance.as_string.return_value ) @patch('argparse.ArgumentParser.parse_args') @patch('blkcore.email_script.get_cusips_data_and_send_email') def test_main(self, mock_get_cusips_data_and_send_email, mock_parse_args): # Simulate command line arguments mock_parse_args.return_value = MagicMock( fileName="cusipList.txt", date="2024-11-11", receiverEmail=["test@example.com"] ) # Call main to trigger the process with patch('sys.stdout', new_callable=StringIO) as mock_stdout: main() # Check if the correct function is called with the expected parameters mock_get_cusips_data_and_send_email.assert_called_once_with( "cusipList.txt", datetime.datetime(2024, 11, 11), ["test@example.com"] ) if __name__ == '__main__': unittest.main()
Editor is loading...
Leave a Comment