Untitled
unknown
plain_text
a year ago
4.8 kB
18
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