"""Test main lambda module.""" import io from os import path from unittest.mock import call from unittest.mock import MagicMock from unittest.mock import patch from zipfile import ZIP_DEFLATED import pandas as pd import pytest import config from src import app from src.app import NumberFormat from src.exceptions import CustomReportException @pytest.mark.parametrize( 'statement_period_names, expected_token', [ ([(2025, 2)], 'feb2025'), ([(2026, 1)], 'jan2026'), ([(2024, 12)], 'dec2024'), ([(2025, 2), (2025, 3)], 'feb-to-mar2025'), ([(2025, 1), (2025, 6)], 'jan-to-jun2025'), ([(2025, 2), (2025, 10)], 'feb-to-oct2025'), ([(2025, 2), (2026, 3)], 'feb2025-to-mar2026'), ([(2024, 12), (2025, 1)], 'dec2024-to-jan2025'), ([(2025, 11), (2026, 2)], 'nov2025-to-feb2026'), ([(2025, 3), (2025, 2)], 'feb-to-mar2025'), ([(2026, 3), (2025, 2)], 'feb2025-to-mar2026'), ([(2025, 6), (2025, 1), (2025, 3)], 'jan-to-jun2025'), ], ) def test_build_statement_period_token(statement_period_names, expected_token): """Test building statement period tokens.""" result = app._build_statement_period_token(statement_period_names) assert result == expected_token @pytest.mark.parametrize( 'report, statement_period_names, expected', [ ( { 'account_id': 123, 'contract_id': None, 'dimension_column': 'one', 'dimension_row': 'two', }, [(2025, 2)], 'report_123_feb2025_one_two.csv', ), ( { 'account_id': 24601, 'contract_id': 10001, 'dimension_column': 'two', 'dimension_row': 'one', }, [(2025, 2), (2025, 10)], 'report_24601_10001_feb-to-oct2025_two_one.csv', ), ( { 'account_id': 123, 'contract_id': None, 'dimension_column': 'one', 'dimension_row': 'service', }, [(2025, 2), (2026, 3)], 'report_123_feb2025-to-mar2026_one_store.csv', ), ( { 'account_id': 123, 'contract_id': None, 'dimension_column': 'one', 'dimension_row': 'two', 'subaccount_id': 12345, }, [(2025, 2)], 'report_12345_feb2025_one_two_subaccount_Test_Subaccount.csv', ), ( { 'account_id': 123, 'contract_id': None, 'dimension_column': 'one', 'dimension_row': 'two', 'subaccount_id': 67890, }, [(2025, 2)], 'report_67890_feb2025_one_two_subaccount_Test_Sub_Account_Name.csv', ), ( { 'account_id': 123456, 'contract_id': 789012, 'dimension_column': 'territory', 'dimension_row': 'product', 'subaccount_id': 11111, }, [(2025, 2), (2025, 3)], 'report_11111_789012_feb-to-mar2025_territory_product_subaccount_Very_Long_Subaccount_Name_That_Will_Definitely_Cause_The_Filename_To_Exceed_Maximum_Character_Limit_For_Filesystem_Compatibility_And_Should_Be_Truncated_Appropriately_To_Fit_Within_The_Al.csv', # noqa: E501 ), ( { 'account_id': 123, 'contract_id': None, 'dimension_column': 'one', 'dimension_row': 'two', 'file_type': 'xls', }, [(2026, 1)], 'report_123_jan2026_one_two.xls', ), ( { 'account_id': 123, 'contract_id': None, 'dimension_column': 'one', 'dimension_row': 'two', 'file_type': 'txt', }, [(2024, 12), (2025, 1)], 'report_123_dec2024-to-jan2025_one_two.txt', ), ( { 'account_id': 123, 'contract_id': None, 'dimension_column': 'one', 'dimension_row': 'two', 'file_type': 'unknown', }, [(2025, 2)], 'report_123_feb2025_one_two.csv', ), ], ) @patch('src.app.snowflake.get_subaccount_name') def test_build_file_name(mock_get_subaccount_name, report, statement_period_names, expected): """Test building a report file name.""" if report.get('subaccount_id') == 12345: mock_get_subaccount_name.return_value = 'Test Subaccount' elif report.get('subaccount_id') == 67890: mock_get_subaccount_name.return_value = 'Test/Sub\\Account Name' elif report.get('subaccount_id') == 11111: mock_get_subaccount_name.return_value = 'Very Long Subaccount Name That Will Definitely Cause The Filename To Exceed Maximum Character Limit For Filesystem Compatibility And Should Be Truncated Appropriately To Fit Within The Allowed Length Of Two Hundred And Fifty Five Characters Maximum For Most File Systems' # noqa: E501 else: mock_get_subaccount_name.return_value = None result = app._build_file_name(report, statement_period_names) assert result == expected @patch('src.app.is_feature_enabled') @patch('src.app._write_file') @patch('src.app._build_file_name') @patch('src.app.OwsMoneyhub') @patch('src.app.s3') @patch('src.app.snowflake') def test_generate_report( snowflake_mock, s3_mock, ows_moneyhub_mock, build_file_name_mock, write_file_mock, is_feature_enabled_mock, report_custom_fixture, ): """Test generating a report with feature flag enabled.""" report_custom_id = 1 filename = 'file.csv' mock_data = MagicMock() ows_moneyhub_mock.get_report_custom.return_value = report_custom_fixture ows_moneyhub_mock.update_report_custom.return_value = report_custom_fixture snowflake_mock.get_report_data.return_value = mock_data snowflake_mock.get_statement_periods_parsed.return_value = [(2025, 2)] build_file_name_mock.return_value = filename is_feature_enabled_mock.return_value = True zip_file = path.join(config.FILE_OUTPUT_PATH, 'file.zip') s3_mock.full_path.side_effect = [zip_file] app._generate_report(report_custom_id) ows_moneyhub_mock.get_report_custom.assert_called_once_with(report_custom_id) snowflake_mock.get_report_data.assert_called_once_with( report_custom_fixture['account_id'], report_custom_fixture['contract_id'], report_custom_fixture['statement_period_ids'], report_custom_fixture['dimension_column'], report_custom_fixture['dimension_row'], report_custom_fixture['revenue_type'], report_custom_fixture['revenue_display_type'], report_custom_fixture.get('subaccount_id'), report_custom_fixture.get('filters'), ) snowflake_mock.get_statement_periods_parsed.assert_called_once_with( report_custom_fixture['statement_period_ids'] ) build_file_name_mock.assert_called_once_with(report_custom_fixture, [(2025, 2)]) is_feature_enabled_mock.assert_called_once_with( 'moneyhub_custom_reports_multiple_files', report_custom_fixture['account_id'] ) write_file_mock.assert_called_once() write_file_mock.assert_called_once_with(zip_file, mock_data, filename, 'csv', NumberFormat.US) s3_mock.upload_file.assert_called_once_with(zip_file, '24601/file.zip') ows_moneyhub_mock.update_report_custom.assert_called_once_with( report_custom_id, file_location=zip_file, report_custom_status=config.CUSTOM_REPORT_STATUS_COMPLETE, ) @patch('src.app.is_feature_enabled') @patch('src.app._write_file_legacy') @patch('src.app._build_file_name') @patch('src.app.OwsMoneyhub') @patch('src.app.s3') @patch('src.app.snowflake') def test_generate_report_legacy_feature_disabled( snowflake_mock, s3_mock, ows_moneyhub_mock, build_file_name_mock, write_file_legacy_mock, is_feature_enabled_mock, report_custom_fixture, ): """Test generating a report with feature flag disabled (legacy mode).""" report_custom_id = 1 filename = 'file.csv' mock_data = MagicMock() ows_moneyhub_mock.get_report_custom.return_value = report_custom_fixture ows_moneyhub_mock.update_report_custom.return_value = report_custom_fixture snowflake_mock.get_report_data.return_value = mock_data snowflake_mock.get_statement_periods_parsed.return_value = [(2025, 2)] build_file_name_mock.return_value = filename is_feature_enabled_mock.return_value = False zip_file = path.join(config.FILE_OUTPUT_PATH, 'file.zip') s3_mock.full_path.side_effect = [zip_file] app._generate_report(report_custom_id) ows_moneyhub_mock.get_report_custom.assert_called_once_with(report_custom_id) snowflake_mock.get_report_data.assert_called_once_with( report_custom_fixture['account_id'], report_custom_fixture['contract_id'], report_custom_fixture['statement_period_ids'], report_custom_fixture['dimension_column'], report_custom_fixture['dimension_row'], report_custom_fixture['revenue_type'], report_custom_fixture['revenue_display_type'], report_custom_fixture.get('subaccount_id'), report_custom_fixture.get('filters'), ) snowflake_mock.get_statement_periods_parsed.assert_called_once_with( report_custom_fixture['statement_period_ids'] ) build_file_name_mock.assert_called_once_with(report_custom_fixture, [(2025, 2)]) is_feature_enabled_mock.assert_called_once_with( 'moneyhub_custom_reports_multiple_files', report_custom_fixture['account_id'] ) write_file_legacy_mock.assert_called_once() write_file_legacy_mock.assert_called_once_with( zip_file, mock_data, filename, 'csv', NumberFormat.US ) s3_mock.upload_file.assert_called_once_with(zip_file, '24601/file.zip') ows_moneyhub_mock.update_report_custom.assert_called_once_with( report_custom_id, file_location=zip_file, report_custom_status=config.CUSTOM_REPORT_STATUS_COMPLETE, ) @patch('src.app.is_feature_enabled') @patch('src.app._write_file') @patch('src.app.OwsMoneyhub') @patch('src.app.snowflake') def test_generate_report_error( snowflake_mock, ows_moneyhub_mock, write_file_mock, is_feature_enabled_mock, report_custom_fixture, ): """Test error handling when generating a report.""" report_custom_id = 1 mock_data = MagicMock() mock_error = CustomReportException('Generic error') ows_moneyhub_mock.get_report_custom.return_value = report_custom_fixture snowflake_mock.get_report_data.return_value = mock_data snowflake_mock.get_statement_periods_parsed.return_value = [(2025, 2)] is_feature_enabled_mock.return_value = True write_file_mock.side_effect = [mock_error] with pytest.raises(CustomReportException) as exception_info: app._generate_report(report_custom_id) assert exception_info.value == mock_error ows_moneyhub_mock.update_report_custom.assert_called_once_with( report_custom_id, file_location=None, report_custom_status=config.CUSTOM_REPORT_STATUS_ERROR ) @patch('src.app.is_feature_enabled') @patch('src.app._write_file') @patch('src.app._build_file_name') @patch('src.app.OwsMoneyhub') @patch('src.app.s3') @patch('src.app.snowflake') def test_generate_report_financial_detail( snowflake_mock, s3_mock, ows_moneyhub_mock, build_file_name_mock, write_file_mock, is_feature_enabled_mock, report_custom_financial_detail_fixture, ): """Test generating a financial detail report.""" report_custom_id = 123 filename = 'file.csv' zip_file = path.join(config.FILE_OUTPUT_PATH, 'file.zip') mock_data = MagicMock() ows_moneyhub_mock.get_report_custom.return_value = report_custom_financial_detail_fixture ows_moneyhub_mock.update_report_custom.return_value = report_custom_financial_detail_fixture snowflake_mock.get_report_data_financial_detail.return_value = mock_data snowflake_mock.get_statement_periods_parsed.return_value = [(2025, 2)] build_file_name_mock.return_value = filename is_feature_enabled_mock.return_value = True mock_file = path.join(config.FILE_OUTPUT_PATH, '{filename}'.format(filename=filename)) s3_mock.full_path.side_effect = [mock_file] app._generate_report(report_custom_id) ows_moneyhub_mock.get_report_custom.assert_called_once_with(report_custom_id) snowflake_mock.get_report_data_financial_detail.assert_called_once_with( report_custom_financial_detail_fixture['account_id'], report_custom_financial_detail_fixture['contract_id'], report_custom_financial_detail_fixture['statement_period_ids'], report_custom_financial_detail_fixture['dimension_row'], report_custom_financial_detail_fixture['revenue_type'], ) snowflake_mock.get_statement_periods_parsed.assert_called_once_with( report_custom_financial_detail_fixture['statement_period_ids'] ) build_file_name_mock.assert_called_once_with( report_custom_financial_detail_fixture, [(2025, 2)] ) write_file_mock.assert_called_once_with(zip_file, mock_data, filename, 'csv', NumberFormat.US) s3_mock.upload_file.assert_called_once_with(zip_file, '24601/file.zip') ows_moneyhub_mock.update_report_custom.assert_called_once_with( report_custom_id, file_location=mock_file, report_custom_status=config.CUSTOM_REPORT_STATUS_COMPLETE, ) @patch('src.app._generate_report') def test_handler(generate_report_mock): """Test handler function.""" report_custom_id = 42 app.handler({'report_custom_id': report_custom_id}, None) generate_report_mock.assert_called_once_with(report_custom_id) @patch('src.app._generate_report') def test_handler_multiple_records(generate_report_mock): """Test handler function.""" data = { 'Records': [ {'body': '{"report_custom_id": 1}'}, {'body': '{"report_custom_id": 2}'}, {'body': '{"report_custom_id": 3}'}, ] } app.handler(data, None) generate_report_mock.assert_has_calls( [ call(1), call(2), call(3), ] ) @patch('src.app.is_feature_enabled') @patch('src.app._write_file') @patch('src.app._build_file_name') @patch('src.app.OwsMoneyhub') @patch('src.app.s3') @patch('src.app.snowflake') def test_generate_report_with_file_type_and_number_format( snowflake_mock, s3_mock, ows_moneyhub_mock, build_file_name_mock, write_file_mock, is_feature_enabled_mock, ): """Test generating a report with specific file type and number format.""" report_custom_id = 1 filename = 'file.xls' zip_file = path.join(config.FILE_OUTPUT_PATH, 'file.zip') report_fixture = { 'report_custom_id': 1, 'account_id': 24601, 'contract_id': None, 'statement_period_ids': '265,266', 'report_custom_status': 'in_progress', 'dimension_column': 'territory', 'dimension_row': 'product', 'revenue_type': config.REVENUE_TYPE_DISTRIBUTION, 'file_location': 's3://file.xls', 'created_by': 'you', 'created_at': '2022-07-25T17:59:12Z', 'file_type': 'xls', 'number_format': 'eu', } mock_data = MagicMock() ows_moneyhub_mock.get_report_custom.return_value = report_fixture ows_moneyhub_mock.update_report_custom.return_value = report_fixture snowflake_mock.get_report_data.return_value = mock_data snowflake_mock.get_statement_periods_parsed.return_value = [(2025, 2)] build_file_name_mock.return_value = filename is_feature_enabled_mock.return_value = True s3_mock.full_path.side_effect = [zip_file] app._generate_report(report_custom_id) snowflake_mock.get_statement_periods_parsed.assert_called_once_with('265,266') build_file_name_mock.assert_called_once_with(report_fixture, [(2025, 2)]) write_file_mock.assert_called_once_with(zip_file, mock_data, filename, 'xls', NumberFormat.EU) def test_make_float_formatter_us(): """Test US number format formatter.""" formatter = app._make_float_formatter(NumberFormat.US) result = formatter(1234.5678901234) assert result == '1234.5678901234' def test_make_float_formatter_eu(): """Test EU number format formatter.""" formatter = app._make_float_formatter(NumberFormat.EU) result = formatter(1234.5678901234) assert result == '1234,5678901234' @patch('src.app.is_feature_enabled') @patch('src.app._write_file') @patch('src.app._build_file_name') @patch('src.app.OwsMoneyhub') @patch('src.app.s3') @patch('src.app.snowflake') def test_generate_report_with_custom_filters( snowflake_mock, s3_mock, ows_moneyhub_mock, build_file_name_mock, write_file_mock, is_feature_enabled_mock, report_custom_fixture, ): """Test generating a report with custom filters.""" report_custom_id = 1 filename = 'file.csv' report_with_filters = report_custom_fixture.copy() report_with_filters['filters'] = { 'artist_ids': [10, 20, 30], 'country_codes': ['US', 'GB'], 'store_ids': [100, 200], } mock_data = MagicMock() ows_moneyhub_mock.get_report_custom.return_value = report_with_filters ows_moneyhub_mock.update_report_custom.return_value = report_with_filters snowflake_mock.get_report_data.return_value = mock_data snowflake_mock.get_statement_periods_parsed.return_value = [(2025, 2)] build_file_name_mock.return_value = filename is_feature_enabled_mock.return_value = True zip_file = path.join(config.FILE_OUTPUT_PATH, 'file.zip') s3_mock.full_path.side_effect = [zip_file] app._generate_report(report_custom_id) expected_filters = { 'artist_ids': [10, 20, 30], 'country_codes': ['US', 'GB'], 'store_ids': [100, 200], } ows_moneyhub_mock.get_report_custom.assert_called_once_with(report_custom_id) snowflake_mock.get_report_data.assert_called_once_with( report_with_filters['account_id'], report_with_filters['contract_id'], report_with_filters['statement_period_ids'], report_with_filters['dimension_column'], report_with_filters['dimension_row'], report_with_filters['revenue_type'], report_with_filters['revenue_display_type'], report_with_filters.get('subaccount_id'), expected_filters, ) snowflake_mock.get_statement_periods_parsed.assert_called_once_with( report_with_filters['statement_period_ids'] ) build_file_name_mock.assert_called_once_with(report_with_filters, [(2025, 2)]) write_file_mock.assert_called_once() s3_mock.upload_file.assert_called_once_with(zip_file, '24601/file.zip') ows_moneyhub_mock.update_report_custom.assert_called_once_with( report_custom_id, file_location=zip_file, report_custom_status=config.CUSTOM_REPORT_STATUS_COMPLETE, ) @patch('src.app.is_feature_enabled') @patch('src.app._write_file') @patch('src.app._build_file_name') @patch('src.app.OwsMoneyhub') @patch('src.app.s3') @patch('src.app.snowflake') def test_generate_report_with_empty_filters( snowflake_mock, s3_mock, ows_moneyhub_mock, build_file_name_mock, write_file_mock, is_feature_enabled_mock, report_custom_fixture, ): """Test generating a report with empty filters.""" report_custom_id = 1 filename = 'file.csv' report_with_filters = report_custom_fixture.copy() report_with_filters['filters'] = {} mock_data = MagicMock() ows_moneyhub_mock.get_report_custom.return_value = report_with_filters ows_moneyhub_mock.update_report_custom.return_value = report_with_filters snowflake_mock.get_report_data.return_value = mock_data snowflake_mock.get_statement_periods_parsed.return_value = [(2025, 2)] build_file_name_mock.return_value = filename is_feature_enabled_mock.return_value = True zip_file = path.join(config.FILE_OUTPUT_PATH, 'file.zip') s3_mock.full_path.side_effect = [zip_file] app._generate_report(report_custom_id) ows_moneyhub_mock.get_report_custom.assert_called_once_with(report_custom_id) snowflake_mock.get_report_data.assert_called_once_with( report_with_filters['account_id'], report_with_filters['contract_id'], report_with_filters['statement_period_ids'], report_with_filters['dimension_column'], report_with_filters['dimension_row'], report_with_filters['revenue_type'], report_with_filters['revenue_display_type'], report_with_filters.get('subaccount_id'), {}, ) snowflake_mock.get_statement_periods_parsed.assert_called_once_with( report_with_filters['statement_period_ids'] ) build_file_name_mock.assert_called_once_with(report_with_filters, [(2025, 2)]) @patch('src.app.pl') def test_write_file_to_buffer_csv(mock_pl): """Test _write_file_to_buffer for CSV file type.""" df = pd.DataFrame({'col1': [1, 2], 'col2': [3.5, 4.5]}) mock_polars_df = MagicMock() mock_pl.from_pandas.return_value = mock_polars_df result = app._write_file_to_buffer(df, 'csv', NumberFormat.US) mock_pl.from_pandas.assert_called_once_with(df, include_index=True) mock_polars_df.write_csv.assert_called_once() call_kwargs = mock_polars_df.write_csv.call_args.kwargs assert call_kwargs['separator'] == ',' assert call_kwargs['quote_style'] == 'non_numeric' assert call_kwargs['decimal_comma'] is False assert call_kwargs['include_bom'] is True assert isinstance(result, io.BytesIO) @patch('src.app.pl') def test_write_file_to_buffer_xls(mock_pl): """Test _write_file_to_buffer for XLS file type.""" df = pd.DataFrame({'col1': [1, 2], 'col2': [3.5, 4.5]}) mock_polars_df = MagicMock() mock_pl.from_pandas.return_value = mock_polars_df app._write_file_to_buffer(df, 'xls', NumberFormat.US) mock_polars_df.write_csv.assert_called_once() call_kwargs = mock_polars_df.write_csv.call_args.kwargs assert call_kwargs['separator'] == '\t' assert call_kwargs['quote_style'] == 'always' assert call_kwargs['include_bom'] is False @patch('src.app._sanitize_for_tsv') @patch('src.app.pl') def test_write_file_to_buffer_txt(mock_pl, mock_sanitize): """Test _write_file_to_buffer for TXT file type.""" df = pd.DataFrame({'col1': [1, 2], 'col2': [3.5, 4.5]}) sanitized_df = pd.DataFrame({'col1': [1, 2], 'col2': [3.5, 4.5]}) mock_sanitize.return_value = sanitized_df mock_polars_df = MagicMock() mock_pl.from_pandas.return_value = mock_polars_df app._write_file_to_buffer(df, 'txt', NumberFormat.US) mock_sanitize.assert_called_once_with(df) mock_pl.from_pandas.assert_called_once_with(sanitized_df, include_index=True) call_kwargs = mock_polars_df.write_csv.call_args.kwargs assert call_kwargs['separator'] == '\t' assert call_kwargs['quote_style'] == 'never' assert call_kwargs['include_bom'] is False @patch('src.app._write_file_to_buffer') @patch('src.app.ZipFile', create=True) def test_write_file_single_file(mock_zipfile, mock_write_buffer): """Test _write_file with dataframe under MAX_ROWS_PER_FILE.""" df = pd.DataFrame({'col1': range(1000), 'col2': range(1000)}) mock_file = MagicMock() mock_zipfile.return_value.__enter__.return_value = mock_file mock_buffer = MagicMock() mock_buffer.getvalue.return_value = b'file contents' mock_write_buffer.return_value = mock_buffer app._write_file('/path/to/file.zip', df, 'report.csv', 'csv', NumberFormat.US) mock_zipfile.assert_called_once_with('/path/to/file.zip', 'w', ZIP_DEFLATED) mock_write_buffer.assert_called_once_with(df, 'csv', NumberFormat.US) mock_file.writestr.assert_called_once_with('report.csv', b'file contents') @patch('src.app._write_file_to_buffer') @patch('src.app.ZipFile', create=True) def test_write_file_legacy(mock_zipfile, mock_write_buffer): """Test _write_file_legacy creates a single file without splitting.""" df = pd.DataFrame({'col1': range(5000), 'col2': range(5000)}) mock_file = MagicMock() mock_zipfile.return_value.__enter__.return_value = mock_file mock_buffer = MagicMock() mock_buffer.getvalue.return_value = b'file contents' mock_write_buffer.return_value = mock_buffer app._write_file_legacy('/path/to/file.zip', df, 'report.csv', 'csv', NumberFormat.US) mock_zipfile.assert_called_once_with('/path/to/file.zip', 'w', ZIP_DEFLATED) mock_write_buffer.assert_called_once_with(df, 'csv', NumberFormat.US) mock_file.writestr.assert_called_once_with('report.csv', b'file contents') @patch('src.app._write_file_to_buffer') @patch('src.app.ZipFile', create=True) @patch('src.app.MAX_ROWS_PER_FILE', 1000) def test_write_file_multiple_files(mock_zipfile, mock_write_buffer): """Test _write_file splits large dataframes into multiple files.""" df = pd.DataFrame({'col1': range(2500), 'col2': range(2500)}) mock_file = MagicMock() mock_zipfile.return_value.__enter__.return_value = mock_file mock_buffer = MagicMock() mock_buffer.getvalue.return_value = b'file contents' mock_write_buffer.return_value = mock_buffer app._write_file('/path/to/file.zip', df, 'report.csv', 'csv', NumberFormat.US) call_args = [call[0] for call in mock_file.writestr.call_args_list] assert mock_write_buffer.call_count == 3 assert mock_file.writestr.call_count == 3 mock_zipfile.assert_called_once_with('/path/to/file.zip', 'w', ZIP_DEFLATED) assert call_args[0][0] == 'report_1.csv' assert call_args[1][0] == 'report_2.csv' assert call_args[2][0] == 'report_3.csv' @patch('src.app._write_file_to_buffer') @patch('src.app.ZipFile', create=True) @patch('src.app.MAX_ROWS_PER_FILE', 1000) def test_write_file_multiple_files_different_extension(mock_zipfile, mock_write_buffer): """Test _write_file splits files with correct extension.""" df = pd.DataFrame({'col1': range(2500), 'col2': range(2500)}) mock_file = MagicMock() mock_zipfile.return_value.__enter__.return_value = mock_file mock_buffer = MagicMock() mock_buffer.getvalue.return_value = b'file contents' mock_write_buffer.return_value = mock_buffer app._write_file('/path/to/file.zip', df, 'report.xls', 'xls', NumberFormat.US) call_args = [call[0] for call in mock_file.writestr.call_args_list] assert call_args[0][0] == 'report_1.xls' assert call_args[1][0] == 'report_2.xls' assert call_args[2][0] == 'report_3.xls' @pytest.mark.parametrize( 'dataframe', [ None, pd.DataFrame(), ], ) @patch('src.app.ZipFile', create=True) def test_write_file_empty_dataframe(mock_zipfile, dataframe): """Test _write_file with empty dataframe.""" mock_file = MagicMock() mock_zipfile.return_value.__enter__.return_value = mock_file app._write_file('/path/to/file.zip', dataframe, '/path/to/file.csv', 'csv', NumberFormat.US) mock_zipfile.assert_called_once_with('/path/to/file.zip', 'w', ZIP_DEFLATED) @pytest.mark.parametrize( 'file_name, file_type, expected_separator, expected_quoting, expected_bom', [ ('file.csv', 'csv', ',', 'non_numeric', True), ('file.xls', 'xls', '\t', 'always', False), ('file.txt', 'txt', '\t', 'never', False), ], ) @patch('src.app._sanitize_for_tsv') @patch('src.app.pl') @patch('src.app._write_file_to_buffer') @patch('src.app.ZipFile', create=True) def test_write_file_different_types( mock_zipfile, mock_write_buffer, mock_pl, mock_sanitize, file_name, file_type, expected_separator, expected_quoting, expected_bom, ): """Test _write_file with different file types and their settings.""" mock_df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]}) mock_file = MagicMock() mock_zipfile.return_value.__enter__.return_value = mock_file mock_buffer = MagicMock() mock_buffer.getvalue.return_value = b'file contents' mock_write_buffer.return_value = mock_buffer app._write_file('/path/to/file.zip', mock_df, file_name, file_type, NumberFormat.US) mock_zipfile.assert_called_with('/path/to/file.zip', 'w', ZIP_DEFLATED) mock_write_buffer.assert_called_once_with(mock_df, file_type, NumberFormat.US) mock_file.writestr.assert_called_once_with(file_name, b'file contents')