Untitled

 avatar
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