# Copyright 2014 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest

import mock
import requests
from six.moves import http_client


def _make_credentials():
    import google.auth.credentials

    return mock.Mock(spec=google.auth.credentials.Credentials)


def _make_response(status=http_client.OK, content=b'', headers={}):
    response = requests.Response()
    response.status_code = status
    response._content = content
    response.headers = headers
    response.request = requests.Request()
    return response


def _make_requests_session(responses):
    session = mock.create_autospec(requests.Session, instance=True)
    session.request.side_effect = responses
    return session


class TestMIMEApplicationHTTP(unittest.TestCase):

    @staticmethod
    def _get_target_class():
        from google.cloud.storage.batch import MIMEApplicationHTTP

        return MIMEApplicationHTTP

    def _make_one(self, *args, **kw):
        return self._get_target_class()(*args, **kw)

    def test_ctor_body_None(self):
        METHOD = 'DELETE'
        PATH = '/path/to/api'
        LINES = [
            "DELETE /path/to/api HTTP/1.1",
            "",
        ]
        mah = self._make_one(METHOD, PATH, {}, None)
        self.assertEqual(mah.get_content_type(), 'application/http')
        self.assertEqual(mah.get_payload().splitlines(), LINES)

    def test_ctor_body_str(self):
        METHOD = 'GET'
        PATH = '/path/to/api'
        BODY = 'ABC'
        HEADERS = {'Content-Length': len(BODY), 'Content-Type': 'text/plain'}
        LINES = [
            "GET /path/to/api HTTP/1.1",
            "Content-Length: 3",
            "Content-Type: text/plain",
            "",
            "ABC",
        ]
        mah = self._make_one(METHOD, PATH, HEADERS, BODY)
        self.assertEqual(mah.get_payload().splitlines(), LINES)

    def test_ctor_body_dict(self):
        METHOD = 'GET'
        PATH = '/path/to/api'
        BODY = {'foo': 'bar'}
        HEADERS = {}
        LINES = [
            'GET /path/to/api HTTP/1.1',
            'Content-Length: 14',
            'Content-Type: application/json',
            '',
            '{"foo": "bar"}',
        ]
        mah = self._make_one(METHOD, PATH, HEADERS, BODY)
        self.assertEqual(mah.get_payload().splitlines(), LINES)


class TestBatch(unittest.TestCase):

    @staticmethod
    def _get_target_class():
        from google.cloud.storage.batch import Batch

        return Batch

    def _make_one(self, *args, **kw):
        return self._get_target_class()(*args, **kw)

    def test_ctor(self):
        http = _make_requests_session([])
        connection = _Connection(http=http)
        client = _Client(connection)
        batch = self._make_one(client)
        self.assertIs(batch._client, client)
        self.assertEqual(len(batch._requests), 0)
        self.assertEqual(len(batch._target_objects), 0)

    def test_current(self):
        from google.cloud.storage.client import Client

        project = 'PROJECT'
        credentials = _make_credentials()
        client = Client(project=project, credentials=credentials)
        batch1 = self._make_one(client)
        self.assertIsNone(batch1.current())

        client._push_batch(batch1)
        self.assertIs(batch1.current(), batch1)

        batch2 = self._make_one(client)
        client._push_batch(batch2)
        self.assertIs(batch1.current(), batch2)

    def test__make_request_GET_normal(self):
        from google.cloud.storage.batch import _FutureDict

        url = 'http://example.com/api'
        http = _make_requests_session([])
        connection = _Connection(http=http)
        batch = self._make_one(connection)
        target = _MockObject()

        response = batch._make_request('GET', url, target_object=target)

        # Check the respone
        self.assertEqual(response.status_code, 204)
        self.assertIsInstance(response.json(), _FutureDict)
        self.assertIsInstance(response.content, _FutureDict)
        self.assertIs(target._properties, response.content)

        # The real http request should not have been called yet.
        http.request.assert_not_called()

        # Check the queued request
        self.assertEqual(len(batch._requests), 1)
        request = batch._requests[0]
        request_method, request_url, _, request_data = request
        self.assertEqual(request_method, 'GET')
        self.assertEqual(request_url, url)
        self.assertIsNone(request_data)

    def test__make_request_POST_normal(self):
        from google.cloud.storage.batch import _FutureDict

        url = 'http://example.com/api'
        http = _make_requests_session([])
        connection = _Connection(http=http)
        batch = self._make_one(connection)
        data = {'foo': 1}
        target = _MockObject()

        response = batch._make_request(
            'POST', url, data={'foo': 1}, target_object=target)

        self.assertEqual(response.status_code, 204)
        self.assertIsInstance(response.content, _FutureDict)
        self.assertIs(target._properties, response.content)

        # The real http request should not have been called yet.
        http.request.assert_not_called()

        request = batch._requests[0]
        request_method, request_url, _, request_data = request
        self.assertEqual(request_method, 'POST')
        self.assertEqual(request_url, url)
        self.assertEqual(request_data, data)

    def test__make_request_PATCH_normal(self):
        from google.cloud.storage.batch import _FutureDict

        url = 'http://example.com/api'
        http = _make_requests_session([])
        connection = _Connection(http=http)
        batch = self._make_one(connection)
        data = {'foo': 1}
        target = _MockObject()

        response = batch._make_request(
            'PATCH', url, data={'foo': 1}, target_object=target)

        self.assertEqual(response.status_code, 204)
        self.assertIsInstance(response.content, _FutureDict)
        self.assertIs(target._properties, response.content)

        # The real http request should not have been called yet.
        http.request.assert_not_called()

        request = batch._requests[0]
        request_method, request_url, _, request_data = request
        self.assertEqual(request_method, 'PATCH')
        self.assertEqual(request_url, url)
        self.assertEqual(request_data, data)

    def test__make_request_DELETE_normal(self):
        from google.cloud.storage.batch import _FutureDict

        url = 'http://example.com/api'
        http = _make_requests_session([])
        connection = _Connection(http=http)
        batch = self._make_one(connection)
        target = _MockObject()

        response = batch._make_request('DELETE', url, target_object=target)

        # Check the respone
        self.assertEqual(response.status_code, 204)
        self.assertIsInstance(response.content, _FutureDict)
        self.assertIs(target._properties, response.content)

        # The real http request should not have been called yet.
        http.request.assert_not_called()

        # Check the queued request
        self.assertEqual(len(batch._requests), 1)
        request = batch._requests[0]
        request_method, request_url, _, request_data = request
        self.assertEqual(request_method, 'DELETE')
        self.assertEqual(request_url, url)
        self.assertIsNone(request_data)

    def test__make_request_POST_too_many_requests(self):
        url = 'http://example.com/api'
        http = _make_requests_session([])
        connection = _Connection(http=http)
        batch = self._make_one(connection)

        batch._MAX_BATCH_SIZE = 1
        batch._requests.append(('POST', url, {}, {'bar': 2}))

        with self.assertRaises(ValueError):
            batch._make_request('POST', url, data={'foo': 1})

    def test_finish_empty(self):
        http = _make_requests_session([])
        connection = _Connection(http=http)
        batch = self._make_one(connection)

        with self.assertRaises(ValueError):
            batch.finish()

    def _get_payload_chunks(self, boundary, payload):
        divider = '--' + boundary[len('boundary="'):-1]
        chunks = payload.split(divider)[1:-1]  # discard prolog / epilog
        return chunks

    def _check_subrequest_no_payload(self, chunk, method, url):
        lines = chunk.splitlines()
        # blank + 2 headers + blank + request + blank + blank
        self.assertEqual(len(lines), 7)
        self.assertEqual(lines[0], '')
        self.assertEqual(lines[1], 'Content-Type: application/http')
        self.assertEqual(lines[2], 'MIME-Version: 1.0')
        self.assertEqual(lines[3], '')
        self.assertEqual(lines[4], '%s %s HTTP/1.1' % (method, url))
        self.assertEqual(lines[5], '')
        self.assertEqual(lines[6], '')

    def _check_subrequest_payload(self, chunk, method, url, payload):
        import json

        lines = chunk.splitlines()
        # blank + 2 headers + blank + request + 2 headers + blank + body
        payload_str = json.dumps(payload)
        self.assertEqual(lines[0], '')
        self.assertEqual(lines[1], 'Content-Type: application/http')
        self.assertEqual(lines[2], 'MIME-Version: 1.0')
        self.assertEqual(lines[3], '')
        self.assertEqual(lines[4], '%s %s HTTP/1.1' % (method, url))
        if method == 'GET':
            self.assertEqual(len(lines), 7)
            self.assertEqual(lines[5], '')
            self.assertEqual(lines[6], '')
        else:
            self.assertEqual(len(lines), 9)
            self.assertEqual(lines[5], 'Content-Length: %d' % len(payload_str))
            self.assertEqual(lines[6], 'Content-Type: application/json')
            self.assertEqual(lines[7], '')
            self.assertEqual(json.loads(lines[8]), payload)

    def _get_mutlipart_request(self, http):
        request_call = http.request.mock_calls[0][2]
        request_headers = request_call['headers']
        request_body = request_call['data']
        content_type, boundary = [
            value.strip() for value in
            request_headers['Content-Type'].split(';')]

        return request_headers, request_body, content_type, boundary

    def test_finish_nonempty(self):
        url = 'http://api.example.com/other_api'
        expected_response = _make_response(
            content=_THREE_PART_MIME_RESPONSE,
            headers={'content-type': 'multipart/mixed; boundary="DEADBEEF="'})
        http = _make_requests_session([expected_response])
        connection = _Connection(http=http)
        client = _Client(connection)
        batch = self._make_one(client)
        batch.API_BASE_URL = 'http://api.example.com'

        batch._do_request('POST', url, {}, {'foo': 1, 'bar': 2}, None)
        batch._do_request('PATCH', url, {}, {'bar': 3}, None)
        batch._do_request('DELETE', url, {}, None, None)
        result = batch.finish()

        self.assertEqual(len(result), len(batch._requests))

        response1, response2, response3 = result

        self.assertEqual(response1.headers, {
            'Content-Length': '20',
            'Content-Type': 'application/json; charset=UTF-8',
        })
        self.assertEqual(response1.json(), {'foo': 1, 'bar': 2})

        self.assertEqual(response2.headers, {
            'Content-Length': '20',
            'Content-Type': 'application/json; charset=UTF-8',
        })
        self.assertEqual(response2.json(), {'foo': 1, 'bar': 3})

        self.assertEqual(response3.headers, {'Content-Length': '0'})
        self.assertEqual(response3.status_code, http_client.NO_CONTENT)

        expected_url = '{}/batch/storage/v1'.format(batch.API_BASE_URL)
        http.request.assert_called_once_with(
            method='POST', url=expected_url, headers=mock.ANY, data=mock.ANY)

        request_info = self._get_mutlipart_request(http)
        request_headers, request_body, content_type, boundary = request_info

        self.assertEqual(content_type, 'multipart/mixed')
        self.assertTrue(boundary.startswith('boundary="=='))
        self.assertTrue(boundary.endswith('=="'))
        self.assertEqual(request_headers['MIME-Version'], '1.0')

        chunks = self._get_payload_chunks(boundary, request_body)
        self.assertEqual(len(chunks), 3)
        self._check_subrequest_payload(
            chunks[0], 'POST', url, {'foo': 1, 'bar': 2})
        self._check_subrequest_payload(chunks[1], 'PATCH', url, {'bar': 3})
        self._check_subrequest_no_payload(chunks[2], 'DELETE', url)

    def test_finish_responses_mismatch(self):
        url = 'http://api.example.com/other_api'
        expected_response = _make_response(
            content=_TWO_PART_MIME_RESPONSE_WITH_FAIL,
            headers={'content-type': 'multipart/mixed; boundary="DEADBEEF="'})
        http = _make_requests_session([expected_response])
        connection = _Connection(http=http)
        client = _Client(connection)
        batch = self._make_one(client)
        batch.API_BASE_URL = 'http://api.example.com'

        batch._requests.append(('GET', url, {}, None))
        with self.assertRaises(ValueError):
            batch.finish()

    def test_finish_nonempty_with_status_failure(self):
        from google.cloud.exceptions import NotFound
        url = 'http://api.example.com/other_api'
        expected_response = _make_response(
            content=_TWO_PART_MIME_RESPONSE_WITH_FAIL,
            headers={'content-type': 'multipart/mixed; boundary="DEADBEEF="'})
        http = _make_requests_session([expected_response])
        connection = _Connection(http=http)
        client = _Client(connection)
        batch = self._make_one(client)
        batch.API_BASE_URL = 'http://api.example.com'
        target1 = _MockObject()
        target2 = _MockObject()

        batch._do_request('GET', url, {}, None, target1)
        batch._do_request('GET', url, {}, None, target2)

        # Make sure futures are not populated.
        self.assertEqual([future for future in batch._target_objects],
                         [target1, target2])
        target2_future_before = target2._properties

        with self.assertRaises(NotFound):
            batch.finish()

        self.assertEqual(target1._properties,
                         {'foo': 1, 'bar': 2})
        self.assertIs(target2._properties, target2_future_before)

        expected_url = '{}/batch/storage/v1'.format(batch.API_BASE_URL)
        http.request.assert_called_once_with(
            method='POST', url=expected_url, headers=mock.ANY, data=mock.ANY)

        _, request_body, _, boundary = self._get_mutlipart_request(http)

        chunks = self._get_payload_chunks(boundary, request_body)
        self.assertEqual(len(chunks), 2)
        self._check_subrequest_payload(chunks[0], 'GET', url, {})
        self._check_subrequest_payload(chunks[1], 'GET', url, {})

    def test_finish_nonempty_non_multipart_response(self):
        url = 'http://api.example.com/other_api'
        http = _make_requests_session([_make_response()])
        connection = _Connection(http=http)
        client = _Client(connection)
        batch = self._make_one(client)
        batch._requests.append(('POST', url, {}, {'foo': 1, 'bar': 2}))

        with self.assertRaises(ValueError):
            batch.finish()

    def test_as_context_mgr_wo_error(self):
        from google.cloud.storage.client import Client

        url = 'http://example.com/api'
        expected_response = _make_response(
            content=_THREE_PART_MIME_RESPONSE,
            headers={'content-type': 'multipart/mixed; boundary="DEADBEEF="'})
        http = _make_requests_session([expected_response])
        project = 'PROJECT'
        credentials = _make_credentials()
        client = Client(project=project, credentials=credentials)
        client._http_internal = http

        self.assertEqual(list(client._batch_stack), [])

        target1 = _MockObject()
        target2 = _MockObject()
        target3 = _MockObject()

        with self._make_one(client) as batch:
            self.assertEqual(list(client._batch_stack), [batch])
            batch._make_request('POST', url, {'foo': 1, 'bar': 2},
                                target_object=target1)
            batch._make_request('PATCH', url, {'bar': 3},
                                target_object=target2)
            batch._make_request('DELETE', url, target_object=target3)

        self.assertEqual(list(client._batch_stack), [])
        self.assertEqual(len(batch._requests), 3)
        self.assertEqual(batch._requests[0][0], 'POST')
        self.assertEqual(batch._requests[1][0], 'PATCH')
        self.assertEqual(batch._requests[2][0], 'DELETE')
        self.assertEqual(batch._target_objects, [target1, target2, target3])
        self.assertEqual(target1._properties,
                         {'foo': 1, 'bar': 2})
        self.assertEqual(target2._properties,
                         {'foo': 1, 'bar': 3})
        self.assertEqual(target3._properties, b'')

    def test_as_context_mgr_w_error(self):
        from google.cloud.storage.batch import _FutureDict
        from google.cloud.storage.client import Client

        URL = 'http://example.com/api'
        http = _make_requests_session([])
        connection = _Connection(http=http)
        project = 'PROJECT'
        credentials = _make_credentials()
        client = Client(project=project, credentials=credentials)
        client._base_connection = connection

        self.assertEqual(list(client._batch_stack), [])

        target1 = _MockObject()
        target2 = _MockObject()
        target3 = _MockObject()
        try:
            with self._make_one(client) as batch:
                self.assertEqual(list(client._batch_stack), [batch])
                batch._make_request('POST', URL, {'foo': 1, 'bar': 2},
                                    target_object=target1)
                batch._make_request('PATCH', URL, {'bar': 3},
                                    target_object=target2)
                batch._make_request('DELETE', URL, target_object=target3)
                raise ValueError()
        except ValueError:
            pass

        http.request.assert_not_called()
        self.assertEqual(list(client._batch_stack), [])
        self.assertEqual(len(batch._requests), 3)
        self.assertEqual(batch._target_objects, [target1, target2, target3])
        # Since the context manager fails, finish will not get called and
        # the _properties will still be futures.
        self.assertIsInstance(target1._properties, _FutureDict)
        self.assertIsInstance(target2._properties, _FutureDict)
        self.assertIsInstance(target3._properties, _FutureDict)


class Test__unpack_batch_response(unittest.TestCase):

    def _call_fut(self, headers, content):
        from google.cloud.storage.batch import _unpack_batch_response

        response = _make_response(content=content, headers=headers)

        return _unpack_batch_response(response)

    def _unpack_helper(self, response, content):
        result = list(self._call_fut(response, content))
        self.assertEqual(len(result), 3)

        self.assertEqual(result[0].status_code, http_client.OK)
        self.assertEqual(result[0].json(), {u'bar': 2, u'foo': 1})
        self.assertEqual(result[1].status_code, http_client.OK)
        self.assertEqual(result[1].json(), {u'foo': 1, u'bar': 3})
        self.assertEqual(result[2].status_code, http_client.NO_CONTENT)

    def test_bytes_headers(self):
        RESPONSE = {'content-type': b'multipart/mixed; boundary="DEADBEEF="'}
        CONTENT = _THREE_PART_MIME_RESPONSE
        self._unpack_helper(RESPONSE, CONTENT)

    def test_unicode_headers(self):
        RESPONSE = {'content-type': u'multipart/mixed; boundary="DEADBEEF="'}
        CONTENT = _THREE_PART_MIME_RESPONSE
        self._unpack_helper(RESPONSE, CONTENT)


_TWO_PART_MIME_RESPONSE_WITH_FAIL = b"""\
--DEADBEEF=
Content-Type: application/json
Content-ID: <response-8a09ca85-8d1d-4f45-9eb0-da8e8b07ec83+1>

HTTP/1.1 200 OK
Content-Type: application/json; charset=UTF-8
Content-Length: 20

{"foo": 1, "bar": 2}

--DEADBEEF=
Content-Type: application/json
Content-ID: <response-8a09ca85-8d1d-4f45-9eb0-da8e8b07ec83+2>

HTTP/1.1 404 Not Found
Content-Type: application/json; charset=UTF-8
Content-Length: 35

{"error": {"message": "Not Found"}}

--DEADBEEF=--
"""

_THREE_PART_MIME_RESPONSE = b"""\
--DEADBEEF=
Content-Type: application/json
Content-ID: <response-8a09ca85-8d1d-4f45-9eb0-da8e8b07ec83+1>

HTTP/1.1 200 OK
Content-Type: application/json; charset=UTF-8
Content-Length: 20

{"foo": 1, "bar": 2}

--DEADBEEF=
Content-Type: application/json
Content-ID: <response-8a09ca85-8d1d-4f45-9eb0-da8e8b07ec83+2>

HTTP/1.1 200 OK
Content-Type: application/json; charset=UTF-8
Content-Length: 20

{"foo": 1, "bar": 3}

--DEADBEEF=
Content-Type: text/plain
Content-ID: <response-8a09ca85-8d1d-4f45-9eb0-da8e8b07ec83+3>

HTTP/1.1 204 No Content
Content-Length: 0

--DEADBEEF=--
"""


class Test__FutureDict(unittest.TestCase):

    def _make_one(self, *args, **kw):
        from google.cloud.storage.batch import _FutureDict

        return _FutureDict(*args, **kw)

    def test_get(self):
        future = self._make_one()
        self.assertRaises(KeyError, future.get, None)

    def test___getitem__(self):
        future = self._make_one()
        value = orig_value = object()
        with self.assertRaises(KeyError):
            value = future[None]
        self.assertIs(value, orig_value)

    def test___setitem__(self):
        future = self._make_one()
        with self.assertRaises(KeyError):
            future[None] = None


class _Connection(object):

    project = 'TESTING'

    def __init__(self, **kw):
        self.__dict__.update(kw)

    def _make_request(self, method, url, data=None, headers=None):
        return self.http.request(url=url, method=method,
                                 headers=headers, data=data)


class _MockObject(object):
    pass


class _Client(object):

    def __init__(self, connection):
        self._base_connection = connection
