rate_test for pydocumentdb

rate_test.py

import unittest
import socket

import json
import pydocumentdb.https_connection
from collections import deque

two_document_response = {
    'Documents': [
        {
            'id' : 1
        },
        {
            'id' : 2
        }
    ]
}

def document_at(response, index):
    """
    return response, but with 'Documents' array truncated to the
    element at position `index`.
    """
    ret = {}
    for k, v in response.items():
        if k == 'Documents':
            v = [v[index]]
        ret[k] = v
    return ret

class MockHTTPResponse:
    version = 11
    def __init__(self, status=200, reason="OK", data="", headers=None):
        self.status = status
        self.reason = reason
        if headers:
            self.headers = dict(headers)
        else:
            self.headers = {}
        self.data = data

    def getheader(name, default=None):
        return self.headers.get(name, default)

    def getheaders(self):
        return self.headers.items()

    def read(self):
        return self.data

class MockHttpsConnection:
    # queue of responses to emit. This variable is shared across all instances
    # (does not support parallel unit tests)
    responses = deque()
    def __init__(self, host, port=None, ssl_configuration=None, strict=None,
            timeout=socket._GLOBAL_DEFAULT_TIMEOUT, source_address=None):
        self.host = host
        self.port = port
        print "created"

    @classmethod
    def add_response(cls, status=200, reason="OK", data="", headers=None):
        """
        Add a  response for all instances of MockHttpsConnection
        """
        MockHttpsConnection.responses.append((status, reason, data, headers))

    def request(self, method, url, body, headers):
        self.method = method
        self.url = url
        self.body = body
        self.headers = headers
        print method, url, body

    def getresponse(self, *args):
        """
        remove the next element from the global responses list and return.
        """
        resp = MockHttpsConnection.responses.popleft()
        print "Retrurning response", resp, "remaining: ", MockHttpsConnection.responses
        hresp = MockHTTPResponse(*resp)
        if len(resp):
            hresp.headers['x-ms-continuation'] = 'yes' # actually, this should be an rid. Works anyway.
        return hresp

pydocumentdb.https_connection.HTTPSConnection = MockHttpsConnection

import pydocumentdb.document_client as document_client
import pydocumentdb.http_constants as http_constants

#masterKey = 'w3r722+hUT6FwzC8MH2gYiqBWUlPg33O4d5FY92qxm1l05gR/IZm+WEOJaiOv0mwAvOPxLuaL5ewHHV7rb/oIw=='
masterKey = ''

class RateTest(unittest.TestCase):

    def setUp(self):
        MockHttpsConnection.responses.clear()
        pass

    def test_retry_after__fail_on_continuation(self):
        print "##### TEST CONT #####"
        MockHttpsConnection.add_response(200, "OK", json.dumps(document_at(two_document_response, 0)))
        MockHttpsConnection.add_response(429, "Too many requests", "{}", {"x-ms-retry-after-ms": 100})
        MockHttpsConnection.add_response(200, "OK", json.dumps(document_at(two_document_response, 1)))

        dc = document_client.DocumentClient("https://localhost:443", {'masterKey' : masterKey })
        it = dc.QueryDocuments('coll_1', "SELECT * FROM coll_1")
        it = iter(it)
        self.assertEqual(1, next(it)['id'])
        self.assertEqual(2, next(it)['id'])
        self.assertEqual(0, len(MockHttpsConnection.responses))

    def test_retry_after__fail_immediately(self):
        print "##### TEST IMMED ######"
        self.assertEqual(0, len(MockHttpsConnection.responses))
        MockHttpsConnection.add_response(429, "Too many requests", "{}", {"x-ms-retry-after-ms": 100})
        MockHttpsConnection.add_response(200, "OK", json.dumps(two_document_response))

        dc = document_client.DocumentClient("https://localhost:443", {'masterKey' : masterKey })
        it = dc.QueryDocuments('coll_1', "SELECT * FROM coll_1")
        it = iter(it)
        self.assertEqual(2, len(MockHttpsConnection.responses))
        self.assertEqual(1, next(it)['id'])
        self.assertEqual(2, next(it)['id'])
        self.assertEqual(0, len(MockHttpsConnection.responses))


if __name__ == '__main__':
    unittest.main()