"""Tests for google connector."""
import os
from unittest import mock

import pytest

from datalytics import config
from datalytics.connectors.Google.Docs import read_sheet


@pytest.mark.parametrize('upload_back', [True, False])
def test_s3_storage_file(upload_back, mocker):
    """Test s3_storage_file context manager."""
    s3_client_mock = mock.MagicMock()
    s3_client_patch = mocker.patch(
        'datalytics.connectors.Google.Docs.read_sheet.boto3.client',
        return_value=s3_client_mock)

    s3_path = 'foo/bar/baz.ext'

    with read_sheet.s3_storage_file(s3_path, upload_back) as file_path:
        s3_client_patch.assert_called_once_with('s3')
        s3_client_mock.download_file.assert_called_once_with(
            config.S3_BUCKET, s3_path, file_path)
        assert os.path.exists(file_path)

    if upload_back:
        s3_client_mock.upload_file.assert_called_once_with(
            file_path, config.S3_BUCKET, s3_path)
    else:
        s3_client_mock.upload_file.assert_not_called()
    assert not os.path.exists(file_path)


@pytest.mark.parametrize('upload_back', [True, False])
def test_s3_storage_file_handles_exception(upload_back, mocker):
    """Test s3_storage_file context manager handles exceptions."""
    s3_client_mock = mock.MagicMock()
    s3_client_patch = mocker.patch(
        'datalytics.connectors.Google.Docs.read_sheet.boto3.client',
        return_value=s3_client_mock)

    s3_path = 'foo/bar/baz.ext'

    with pytest.raises(Exception):
        with read_sheet.s3_storage_file(s3_path, upload_back) as file_path:
            s3_client_patch.assert_called_once_with('s3')
            s3_client_mock.download_file.assert_called_once_with(
                config.S3_BUCKET, s3_path, file_path)
            assert os.path.exists(file_path)
            raise Exception('Boom!')

    s3_client_mock.upload_file.assert_not_called()
    assert not os.path.exists(file_path)