Repository: home-assistant/core
File: tests/components/aws_s3/test_backup.py (Python)

Code before:
        )

        assert resp.status == 201

        tar_filename, metadata_filename = suggested_filenames(test_backup)

        # Get prefix from config entry
        prefix = mock_config_entry.data.get(CONF_PREFIX)
        if prefix is not None:
            expected_tar_key = f"{prefix}/{tar_filename}"
            expected_metadata_key = f"{prefix}/{metadata_filename}"
        else:
            expected_tar_key = tar_filename
            expected_metadata_key = metadata_filename

        if test_backup.size < MULTIPART_MIN_PART_SIZE_BYTES:
            mock_client.put_object.assert_any_call(
                Bucket="test", Key=expected_tar_key, Body=ANY
            )
            mock_client.put_object.assert_any_call(
                Bucket="test", Key=expected_metadata_key, Body=ANY
            )
        else:
            mock_client.create_multipart_upload.assert_called_with(
                Bucket="test", Key=expected_tar_key
            )
            mock_client.upload_part.assert_any_call(
                Bucket="test",
                Key=expected_tar_key,
                PartNumber=1,
                UploadId="upload_id",
                Body=ANY,
            )
            mock_client.complete_multipart_upload.assert_called_with(
                Bucket="test",
                Key=expected_tar_key,
                UploadId="upload_id",
                MultipartUpload=ANY,
            )
            mock_client.put_object.assert_called_with(
                Bucket="test", Key=expected_metadata_key, Body=ANY
            )


@pytest.mark.parametrize(
    "config_entry_extra_data",
    [
        {"prefix": "backups/home"},
        {},
    ],
    ids=["with_prefix", "no_prefix"],

Review comment (suggestion):
Instead of the branching, can the expected calls be passed in as parametrized call lists?
I think that would require getting rid of the implicit parametrization in `test_backup` though, I think it's OK to do that in a follow-up PR.

Code after:
            mock_client.put_object.assert_any_call(
                Bucket="test", Key=expected_tar_key, Body=ANY
            )
            mock_client.put_object.assert_any_call(
                Bucket="test", Key=expected_metadata_key, Body=ANY
            )
        else:
            mock_client.create_multipart_upload.assert_called_with(
                Bucket="test", Key=expected_tar_key
            )
            mock_client.upload_part.assert_any_call(
                Bucket="test",
                Key=expected_tar_key,
                PartNumber=1,
                UploadId="upload_id",
                Body=ANY,
            )
            mock_client.complete_multipart_upload.assert_called_with(
                Bucket="test",
                Key=expected_tar_key,
                UploadId="upload_id",
                MultipartUpload=ANY,
            )
            mock_client.put_object.assert_called_with(
                Bucket="test", Key=expected_metadata_key, Body=ANY
            )


@pytest.mark.parametrize(
    ("config_entry_extra_data", "expected_key_prefix"),
    [
        ({"prefix": "backups/home"}, "backups/home/"),
        ({}, ""),
    ],
    ids=["with_prefix", "no_prefix"],
)
async def test_agent_download_backup_parametrized(
    hass: HomeAssistant,
    hass_client: ClientSessionGenerator,
    mock_client: MagicMock,
    mock_config_entry: MockConfigEntry,
    test_backup: AgentBackup,
    expected_key_prefix: str,
) -> None:
    """Test agent download backup with and without prefix."""
    client = await hass_client()
    backup_id = "23e64aec"

    resp = await client.get(
        f"/api/backup/download/{backup_id}?agent_id={DOMAIN}.{mock_config_entry.entry_id}"
    )