/y/google-cloud-sdk/platform/gsutil/third_party/boto/tests/unit/dynamodb2/test_table.py

Large files files are truncated, but you can click here to view the full file

from tests.compat import mock, unittest
from boto.dynamodb2 import exceptions
from boto.dynamodb2.fields import (HashKey, RangeKey,
                                   AllIndex, KeysOnlyIndex, IncludeIndex,
                                   GlobalAllIndex, GlobalKeysOnlyIndex,
                                   GlobalIncludeIndex)
from boto.dynamodb2.items import Item
from boto.dynamodb2.layer1 import DynamoDBConnection
from boto.dynamodb2.results import ResultSet, BatchGetResultSet
from boto.dynamodb2.table import Table
from boto.dynamodb2.types import (STRING, NUMBER, BINARY,
                                  FILTER_OPERATORS, QUERY_OPERATORS)
from boto.exception import JSONResponseError
from boto.compat import six, long_type


FakeDynamoDBConnection = mock.create_autospec(DynamoDBConnection)


class SchemaFieldsTestCase(unittest.TestCase):
    def test_hash_key(self):
        hash_key = HashKey('hello')
        self.assertEqual(hash_key.name, 'hello')
        self.assertEqual(hash_key.data_type, STRING)
        self.assertEqual(hash_key.attr_type, 'HASH')

        self.assertEqual(hash_key.definition(), {
            'AttributeName': 'hello',
            'AttributeType': 'S'
        })
        self.assertEqual(hash_key.schema(), {
            'AttributeName': 'hello',
            'KeyType': 'HASH'
        })

    def test_range_key(self):
        range_key = RangeKey('hello')
        self.assertEqual(range_key.name, 'hello')
        self.assertEqual(range_key.data_type, STRING)
        self.assertEqual(range_key.attr_type, 'RANGE')

        self.assertEqual(range_key.definition(), {
            'AttributeName': 'hello',
            'AttributeType': 'S'
        })
        self.assertEqual(range_key.schema(), {
            'AttributeName': 'hello',
            'KeyType': 'RANGE'
        })

    def test_alternate_type(self):
        alt_key = HashKey('alt', data_type=NUMBER)
        self.assertEqual(alt_key.name, 'alt')
        self.assertEqual(alt_key.data_type, NUMBER)
        self.assertEqual(alt_key.attr_type, 'HASH')

        self.assertEqual(alt_key.definition(), {
            'AttributeName': 'alt',
            'AttributeType': 'N'
        })
        self.assertEqual(alt_key.schema(), {
            'AttributeName': 'alt',
            'KeyType': 'HASH'
        })


class IndexFieldTestCase(unittest.TestCase):
    def test_all_index(self):
        all_index = AllIndex('AllKeys', parts=[
            HashKey('username'),
            RangeKey('date_joined')
        ])
        self.assertEqual(all_index.name, 'AllKeys')
        self.assertEqual([part.attr_type for part in all_index.parts], [
            'HASH',
            'RANGE'
        ])
        self.assertEqual(all_index.projection_type, 'ALL')

        self.assertEqual(all_index.definition(), [
            {'AttributeName': 'username', 'AttributeType': 'S'},
            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
        ])
        self.assertEqual(all_index.schema(), {
            'IndexName': 'AllKeys',
            'KeySchema': [
                {
                    'AttributeName': 'username',
                    'KeyType': 'HASH'
                },
                {
                    'AttributeName': 'date_joined',
                    'KeyType': 'RANGE'
                }
            ],
            'Projection': {
                'ProjectionType': 'ALL'
            }
        })

    def test_keys_only_index(self):
        keys_only = KeysOnlyIndex('KeysOnly', parts=[
            HashKey('username'),
            RangeKey('date_joined')
        ])
        self.assertEqual(keys_only.name, 'KeysOnly')
        self.assertEqual([part.attr_type for part in keys_only.parts], [
            'HASH',
            'RANGE'
        ])
        self.assertEqual(keys_only.projection_type, 'KEYS_ONLY')

        self.assertEqual(keys_only.definition(), [
            {'AttributeName': 'username', 'AttributeType': 'S'},
            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
        ])
        self.assertEqual(keys_only.schema(), {
            'IndexName': 'KeysOnly',
            'KeySchema': [
                {
                    'AttributeName': 'username',
                    'KeyType': 'HASH'
                },
                {
                    'AttributeName': 'date_joined',
                    'KeyType': 'RANGE'
                }
            ],
            'Projection': {
                'ProjectionType': 'KEYS_ONLY'
            }
        })

    def test_include_index(self):
        include_index = IncludeIndex('IncludeKeys', parts=[
            HashKey('username'),
            RangeKey('date_joined')
        ], includes=[
            'gender',
            'friend_count'
        ])
        self.assertEqual(include_index.name, 'IncludeKeys')
        self.assertEqual([part.attr_type for part in include_index.parts], [
            'HASH',
            'RANGE'
        ])
        self.assertEqual(include_index.projection_type, 'INCLUDE')

        self.assertEqual(include_index.definition(), [
            {'AttributeName': 'username', 'AttributeType': 'S'},
            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
        ])
        self.assertEqual(include_index.schema(), {
            'IndexName': 'IncludeKeys',
            'KeySchema': [
                {
                    'AttributeName': 'username',
                    'KeyType': 'HASH'
                },
                {
                    'AttributeName': 'date_joined',
                    'KeyType': 'RANGE'
                }
            ],
            'Projection': {
                'ProjectionType': 'INCLUDE',
                'NonKeyAttributes': [
                    'gender',
                    'friend_count',
                ]
            }
        })

    def test_global_all_index(self):
        all_index = GlobalAllIndex('AllKeys', parts=[
            HashKey('username'),
            RangeKey('date_joined')
        ],
        throughput={
            'read': 6,
            'write': 2,
        })
        self.assertEqual(all_index.name, 'AllKeys')
        self.assertEqual([part.attr_type for part in all_index.parts], [
            'HASH',
            'RANGE'
        ])
        self.assertEqual(all_index.projection_type, 'ALL')

        self.assertEqual(all_index.definition(), [
            {'AttributeName': 'username', 'AttributeType': 'S'},
            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
        ])
        self.assertEqual(all_index.schema(), {
            'IndexName': 'AllKeys',
            'KeySchema': [
                {
                    'AttributeName': 'username',
                    'KeyType': 'HASH'
                },
                {
                    'AttributeName': 'date_joined',
                    'KeyType': 'RANGE'
                }
            ],
            'Projection': {
                'ProjectionType': 'ALL'
            },
            'ProvisionedThroughput': {
                'ReadCapacityUnits': 6,
                'WriteCapacityUnits': 2
            }
        })

    def test_global_keys_only_index(self):
        keys_only = GlobalKeysOnlyIndex('KeysOnly', parts=[
            HashKey('username'),
            RangeKey('date_joined')
        ],
        throughput={
            'read': 3,
            'write': 4,
        })
        self.assertEqual(keys_only.name, 'KeysOnly')
        self.assertEqual([part.attr_type for part in keys_only.parts], [
            'HASH',
            'RANGE'
        ])
        self.assertEqual(keys_only.projection_type, 'KEYS_ONLY')

        self.assertEqual(keys_only.definition(), [
            {'AttributeName': 'username', 'AttributeType': 'S'},
            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
        ])
        self.assertEqual(keys_only.schema(), {
            'IndexName': 'KeysOnly',
            'KeySchema': [
                {
                    'AttributeName': 'username',
                    'KeyType': 'HASH'
                },
                {
                    'AttributeName': 'date_joined',
                    'KeyType': 'RANGE'
                }
            ],
            'Projection': {
                'ProjectionType': 'KEYS_ONLY'
            },
            'ProvisionedThroughput': {
                'ReadCapacityUnits': 3,
                'WriteCapacityUnits': 4
            }
        })

    def test_global_include_index(self):
        # Lean on the default throughput
        include_index = GlobalIncludeIndex('IncludeKeys', parts=[
            HashKey('username'),
            RangeKey('date_joined')
        ], includes=[
            'gender',
            'friend_count'
        ])
        self.assertEqual(include_index.name, 'IncludeKeys')
        self.assertEqual([part.attr_type for part in include_index.parts], [
            'HASH',
            'RANGE'
        ])
        self.assertEqual(include_index.projection_type, 'INCLUDE')

        self.assertEqual(include_index.definition(), [
            {'AttributeName': 'username', 'AttributeType': 'S'},
            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
        ])
        self.assertEqual(include_index.schema(), {
            'IndexName': 'IncludeKeys',
            'KeySchema': [
                {
                    'AttributeName': 'username',
                    'KeyType': 'HASH'
                },
                {
                    'AttributeName': 'date_joined',
                    'KeyType': 'RANGE'
                }
            ],
            'Projection': {
                'ProjectionType': 'INCLUDE',
                'NonKeyAttributes': [
                    'gender',
                    'friend_count',
                ]
            },
            'ProvisionedThroughput': {
                'ReadCapacityUnits': 5,
                'WriteCapacityUnits': 5
            }
        })

    def test_global_include_index_throughput(self):
        include_index = GlobalIncludeIndex('IncludeKeys', parts=[
            HashKey('username'),
            RangeKey('date_joined')
        ], includes=[
            'gender',
            'friend_count'
        ], throughput={
            'read': 10,
            'write': 8
        })

        self.assertEqual(include_index.schema(), {
            'IndexName': 'IncludeKeys',
            'KeySchema': [
                {
                    'AttributeName': 'username',
                    'KeyType': 'HASH'
                },
                {
                    'AttributeName': 'date_joined',
                    'KeyType': 'RANGE'
                }
            ],
            'Projection': {
                'ProjectionType': 'INCLUDE',
                'NonKeyAttributes': [
                    'gender',
                    'friend_count',
                ]
            },
            'ProvisionedThroughput': {
                'ReadCapacityUnits': 10,
                'WriteCapacityUnits': 8
            }
        })


class ItemTestCase(unittest.TestCase):
    if six.PY2:
        assertCountEqual = unittest.TestCase.assertItemsEqual

    def setUp(self):
        super(ItemTestCase, self).setUp()
        self.table = Table('whatever', connection=FakeDynamoDBConnection())
        self.johndoe = self.create_item({
            'username': 'johndoe',
            'first_name': 'John',
            'date_joined': 12345,
        })

    def create_item(self, data):
        return Item(self.table, data=data)

    def test_initialization(self):
        empty_item = Item(self.table)
        self.assertEqual(empty_item.table, self.table)
        self.assertEqual(empty_item._data, {})

        full_item = Item(self.table, data={
            'username': 'johndoe',
            'date_joined': 12345,
        })
        self.assertEqual(full_item.table, self.table)
        self.assertEqual(full_item._data, {
            'username': 'johndoe',
            'date_joined': 12345,
        })

    # The next couple methods make use of ``sorted(...)`` so we get consistent
    # ordering everywhere & no erroneous failures.

    def test_keys(self):
        self.assertCountEqual(self.johndoe.keys(), [
            'date_joined',
            'first_name',
            'username',
        ])

    def test_values(self):
        self.assertCountEqual(self.johndoe.values(),
                              [12345, 'John', 'johndoe'])

    def test_contains(self):
        self.assertIn('username', self.johndoe)
        self.assertIn('first_name', self.johndoe)
        self.assertIn('date_joined', self.johndoe)
        self.assertNotIn('whatever', self.johndoe)

    def test_iter(self):
        self.assertCountEqual(self.johndoe,
                              ['johndoe', 'John', 12345])

    def test_get(self):
        self.assertEqual(self.johndoe.get('username'), 'johndoe')
        self.assertEqual(self.johndoe.get('first_name'), 'John')
        self.assertEqual(self.johndoe.get('date_joined'), 12345)

        # Test a missing key. No default yields ``None``.
        self.assertEqual(self.johndoe.get('last_name'), None)
        # This time with a default.
        self.assertEqual(self.johndoe.get('last_name', True), True)

    def test_items(self):
        self.assertCountEqual(
            self.johndoe.items(),
            [
                ('date_joined', 12345),
                ('first_name', 'John'),
                ('username', 'johndoe'),
            ])

    def test_attribute_access(self):
        self.assertEqual(self.johndoe['username'], 'johndoe')
        self.assertEqual(self.johndoe['first_name'], 'John')
        self.assertEqual(self.johndoe['date_joined'], 12345)

        # Test a missing key.
        self.assertEqual(self.johndoe['last_name'], None)

        # Set a key.
        self.johndoe['last_name'] = 'Doe'
        # Test accessing the new key.
        self.assertEqual(self.johndoe['last_name'], 'Doe')

        # Delete a key.
        del self.johndoe['last_name']
        # Test the now-missing-again key.
        self.assertEqual(self.johndoe['last_name'], None)

    def test_needs_save(self):
        self.johndoe.mark_clean()
        self.assertFalse(self.johndoe.needs_save())
        self.johndoe['last_name'] = 'Doe'
        self.assertTrue(self.johndoe.needs_save())

    def test_needs_save_set_changed(self):
        # First, ensure we're clean.
        self.johndoe.mark_clean()
        self.assertFalse(self.johndoe.needs_save())
        # Add a friends collection.
        self.johndoe['friends'] = set(['jane', 'alice'])
        self.assertTrue(self.johndoe.needs_save())
        # Now mark it clean, then change the collection.
        # This does NOT call ``__setitem__``, so the item used to be
        # incorrectly appearing to be clean, when it had in fact been changed.
        self.johndoe.mark_clean()
        self.assertFalse(self.johndoe.needs_save())
        self.johndoe['friends'].add('bob')
        self.assertTrue(self.johndoe.needs_save())

    def test_mark_clean(self):
        self.johndoe['last_name'] = 'Doe'
        self.assertTrue(self.johndoe.needs_save())
        self.johndoe.mark_clean()
        self.assertFalse(self.johndoe.needs_save())

    def test_load(self):
        empty_item = Item(self.table)
        empty_item.load({
            'Item': {
                'username': {'S': 'johndoe'},
                'first_name': {'S': 'John'},
                'last_name': {'S': 'Doe'},
                'date_joined': {'N': '1366056668'},
                'friend_count': {'N': '3'},
                'friends': {'SS': ['alice', 'bob', 'jane']},
            }
        })
        self.assertEqual(empty_item['username'], 'johndoe')
        self.assertEqual(empty_item['date_joined'], 1366056668)
        self.assertEqual(sorted(empty_item['friends']), sorted([
            'alice',
            'bob',
            'jane'
        ]))

    def test_get_keys(self):
        # Setup the data.
        self.table.schema = [
            HashKey('username'),
            RangeKey('date_joined'),
        ]
        self.assertEqual(self.johndoe.get_keys(), {
            'username': 'johndoe',
            'date_joined': 12345,
        })

    def test_get_raw_keys(self):
        # Setup the data.
        self.table.schema = [
            HashKey('username'),
            RangeKey('date_joined'),
        ]
        self.assertEqual(self.johndoe.get_raw_keys(), {
            'username': {'S': 'johndoe'},
            'date_joined': {'N': '12345'},
        })

    def test_build_expects(self):
        # Pristine.
        self.assertEqual(self.johndoe.build_expects(), {
            'first_name': {
                'Exists': False,
            },
            'username': {
                'Exists': False,
            },
            'date_joined': {
                'Exists': False,
            },
        })

        # Without modifications.
        self.johndoe.mark_clean()
        self.assertEqual(self.johndoe.build_expects(), {
            'first_name': {
                'Exists': True,
                'Value': {
                    'S': 'John',
                },
            },
            'username': {
                'Exists': True,
                'Value': {
                    'S': 'johndoe',
                },
            },
            'date_joined': {
                'Exists': True,
                'Value': {
                    'N': '12345',
                },
            },
        })

        # Change some data.
        self.johndoe['first_name'] = 'Johann'
        # Add some data.
        self.johndoe['last_name'] = 'Doe'
        # Delete some data.
        del self.johndoe['date_joined']

        # All fields (default).
        self.assertEqual(self.johndoe.build_expects(), {
            'first_name': {
                'Exists': True,
                'Value': {
                    'S': 'John',
                },
            },
            'last_name': {
                'Exists': False,
            },
            'username': {
                'Exists': True,
                'Value': {
                    'S': 'johndoe',
                },
            },
            'date_joined': {
                'Exists': True,
                'Value': {
                    'N': '12345',
                },
            },
        })

        # Only a subset of the fields.
        self.assertEqual(self.johndoe.build_expects(fields=[
            'first_name',
            'last_name',
            'date_joined',
        ]), {
            'first_name': {
                'Exists': True,
                'Value': {
                    'S': 'John',
                },
            },
            'last_name': {
                'Exists': False,
            },
            'date_joined': {
                'Exists': True,
                'Value': {
                    'N': '12345',
                },
            },
        })

    def test_prepare_full(self):
        self.assertEqual(self.johndoe.prepare_full(), {
            'username': {'S': 'johndoe'},
            'first_name': {'S': 'John'},
            'date_joined': {'N': '12345'}
        })

        self.johndoe['friends'] = set(['jane', 'alice'])
        data = self.johndoe.prepare_full()
        self.assertEqual(data['username'], {'S': 'johndoe'})
        self.assertEqual(data['first_name'], {'S': 'John'})
        self.assertEqual(data['date_joined'], {'N': '12345'})
        self.assertCountEqual(data['friends']['SS'],
                              ['jane', 'alice'])

    def test_prepare_full_empty_set(self):
        self.johndoe['friends'] = set()
        self.assertEqual(self.johndoe.prepare_full(), {
            'username': {'S': 'johndoe'},
            'first_name': {'S': 'John'},
            'date_joined': {'N': '12345'}
        })

    def test_prepare_partial(self):
        self.johndoe.mark_clean()
        # Change some data.
        self.johndoe['first_name'] = 'Johann'
        # Add some data.
        self.johndoe['last_name'] = 'Doe'
        # Delete some data.
        del self.johndoe['date_joined']

        final_data, fields = self.johndoe.prepare_partial()
        self.assertEqual(final_data, {
            'date_joined': {
                'Action': 'DELETE',
            },
            'first_name': {
                'Action': 'PUT',
                'Value': {'S': 'Johann'},
            },
            'last_name': {
                'Action': 'PUT',
                'Value': {'S': 'Doe'},
            },
        })
        self.assertEqual(fields, set([
            'first_name',
            'last_name',
            'date_joined'
        ]))

    def test_prepare_partial_empty_set(self):
        self.johndoe.mark_clean()
        # Change some data.
        self.johndoe['first_name'] = 'Johann'
        # Add some data.
        self.johndoe['last_name'] = 'Doe'
        # Delete some data.
        del self.johndoe['date_joined']
        # Put an empty set on the ``Item``.
        self.johndoe['friends'] = set()

        final_data, fields = self.johndoe.prepare_partial()
        self.assertEqual(final_data, {
            'date_joined': {
                'Action': 'DELETE',
            },
            'first_name': {
                'Action': 'PUT',
                'Value': {'S': 'Johann'},
            },
            'last_name': {
                'Action': 'PUT',
                'Value': {'S': 'Doe'},
            },
        })
        self.assertEqual(fields, set([
            'first_name',
            'last_name',
            'date_joined'
        ]))

    def test_save_no_changes(self):
        # Unchanged, no save.
        with mock.patch.object(self.table, '_put_item', return_value=True) \
                as mock_put_item:
            # Pretend we loaded it via ``get_item``...
            self.johndoe.mark_clean()
            self.assertFalse(self.johndoe.save())

        self.assertFalse(mock_put_item.called)

    def test_save_with_changes(self):
        # With changed data.
        with mock.patch.object(self.table, '_put_item', return_value=True) \
                as mock_put_item:
            self.johndoe.mark_clean()
            self.johndoe['first_name'] = 'J'
            self.johndoe['new_attr'] = 'never_seen_before'
            self.assertTrue(self.johndoe.save())
            self.assertFalse(self.johndoe.needs_save())

        self.assertTrue(mock_put_item.called)
        mock_put_item.assert_called_once_with({
            'username': {'S': 'johndoe'},
            'first_name': {'S': 'J'},
            'new_attr': {'S': 'never_seen_before'},
            'date_joined': {'N': '12345'}
        }, expects={
            'username': {
                'Value': {
                    'S': 'johndoe',
                },
                'Exists': True,
            },
            'first_name': {
                'Value': {
                    'S': 'John',
                },
                'Exists': True,
            },
            'new_attr': {
                'Exists': False,
            },
            'date_joined': {
                'Value': {
                    'N': '12345',
                },
                'Exists': True,
            },
        })

    def test_save_with_changes_overwrite(self):
        # With changed data.
        with mock.patch.object(self.table, '_put_item', return_value=True) \
                as mock_put_item:
            self.johndoe['first_name'] = 'J'
            self.johndoe['new_attr'] = 'never_seen_before'
            # OVERWRITE ALL THE THINGS
            self.assertTrue(self.johndoe.save(overwrite=True))
            self.assertFalse(self.johndoe.needs_save())

        self.assertTrue(mock_put_item.called)
        mock_put_item.assert_called_once_with({
            'username': {'S': 'johndoe'},
            'first_name': {'S': 'J'},
            'new_attr': {'S': 'never_seen_before'},
            'date_joined': {'N': '12345'}
        }, expects=None)

    def test_partial_no_changes(self):
        # Unchanged, no save.
        with mock.patch.object(self.table, '_update_item', return_value=True) \
                as mock_update_item:
            # Pretend we loaded it via ``get_item``...
            self.johndoe.mark_clean()
            self.assertFalse(self.johndoe.partial_save())

        self.assertFalse(mock_update_item.called)

    def test_partial_with_changes(self):
        # Setup the data.
        self.table.schema = [
            HashKey('username'),
        ]

        # With changed data.
        with mock.patch.object(self.table, '_update_item', return_value=True) \
                as mock_update_item:
            # Pretend we loaded it via ``get_item``...
            self.johndoe.mark_clean()
            # Now... MODIFY!!!
            self.johndoe['first_name'] = 'J'
            self.johndoe['last_name'] = 'Doe'
            del self.johndoe['date_joined']
            self.assertTrue(self.johndoe.partial_save())
            self.assertFalse(self.johndoe.needs_save())

        self.assertTrue(mock_update_item.called)
        mock_update_item.assert_called_once_with({
            'username': 'johndoe',
        }, {
            'first_name': {
                'Action': 'PUT',
                'Value': {'S': 'J'},
            },
            'last_name': {
                'Action': 'PUT',
                'Value': {'S': 'Doe'},
            },
            'date_joined': {
                'Action': 'DELETE',
            }
        }, expects={
            'first_name': {
                'Value': {
                    'S': 'John',
                },
                'Exists': True
            },
            'last_name': {
                'Exists': False
            },
            'date_joined': {
                'Value': {
                    'N': '12345',
                },
                'Exists': True
            },
        })

    def test_delete(self):
        # Setup the data.
        self.table.schema = [
            HashKey('username'),
            RangeKey('date_joined'),
        ]

        with mock.patch.object(self.table, 'delete_item', return_value=True) \
                as mock_delete_item:
            self.johndoe.delete()

        self.assertTrue(mock_delete_item.called)
        mock_delete_item.assert_called_once_with(
            username='johndoe',
            date_joined=12345
        )

    def test_nonzero(self):
        self.assertTrue(self.johndoe)
        self.assertFalse(self.create_item({}))


class ItemFromItemTestCase(ItemTestCase):
    def setUp(self):
        super(ItemFromItemTestCase, self).setUp()
        self.johndoe = self.create_item(self.johndoe)


def fake_results(name, greeting='hello', exclusive_start_key=None, limit=None):
    if exclusive_start_key is None:
        exclusive_start_key = -1

    if limit == 0:
        raise Exception("Web Service Returns '400 Bad Request'")

    end_cap = 13
    results = []
    start_key = exclusive_start_key + 1

    for i in range(start_key, start_key + 5):
        if i < end_cap:
            results.append("%s %s #%s" % (greeting, name, i))

    # Don't return more than limit results
    if limit < len(results):
        results = results[:limit]

    retval = {
        'results': results,
    }

    if exclusive_start_key + 5 < end_cap:
        retval['last_key'] = exclusive_start_key + 5

    return retval


class ResultSetTestCase(unittest.TestCase):
    def setUp(self):
        super(ResultSetTestCase, self).setUp()
        self.results = ResultSet()
        self.result_function = mock.MagicMock(side_effect=fake_results)
        self.results.to_call(self.result_function, 'john', greeting='Hello', limit=20)

    def test_first_key(self):
        self.assertEqual(self.results.first_key, 'exclusive_start_key')

    def test_max_page_size_fetch_more(self):
        self.results = ResultSet(max_page_size=10)
        self.results.to_call(self.result_function, 'john', greeting='Hello')
        self.results.fetch_more()
        self.result_function.assert_called_with('john', greeting='Hello', limit=10)
        self.result_function.reset_mock()

    def test_max_page_size_and_smaller_limit_fetch_more(self):
        self.results = ResultSet(max_page_size=10)
        self.results.to_call(self.result_function, 'john', greeting='Hello', limit=5)
        self.results.fetch_more()
        self.result_function.assert_called_with('john', greeting='Hello', limit=5)
        self.result_function.reset_mock()

    def test_max_page_size_and_bigger_limit_fetch_more(self):
        self.results = ResultSet(max_page_size=10)
        self.results.to_call(self.result_function, 'john', greeting='Hello', limit=15)
        self.results.fetch_more()
        self.result_function.assert_called_with('john', greeting='Hello', limit=10)
        self.result_function.reset_mock()

    def test_fetch_more(self):
        # First "page".
        self.results.fetch_more()
        self.assertEqual(self.results._results, [
            'Hello john #0',
            'Hello john #1',
            'Hello john #2',
            'Hello john #3',
            'Hello john #4',
        ])

        self.result_function.assert_called_with('john', greeting='Hello', limit=20)
        self.result_function.reset_mock()

        # Fake in a last key.
        self.results._last_key_seen = 4
        # Second "page".
        self.results.fetch_more()
        self.assertEqual(self.results._results, [
            'Hello john #5',
            'Hello john #6',
            'Hello john #7',
            'Hello john #8',
            'Hello john #9',
        ])

        self.result_function.assert_called_with('john', greeting='Hello', limit=20, exclusive_start_key=4)
        self.result_function.reset_mock()

        # Fake in a last key.
        self.results._last_key_seen = 9
        # Last "page".
        self.results.fetch_more()
        self.assertEqual(self.results._results, [
            'Hello john #10',
            'Hello john #11',
            'Hello john #12',
        ])

        # Fake in a key outside the range.
        self.results._last_key_seen = 15
        # Empty "page". Nothing new gets added
        self.results.fetch_more()
        self.assertEqual(self.results._results, [])

        # Make sure we won't check for results in the future.
        self.assertFalse(self.results._results_left)

    def test_iteration(self):
        # First page.
        self.assertEqual(next(self.results), 'Hello john #0')
        self.assertEqual(next(self.results), 'Hello john #1')
        self.assertEqual(next(self.results), 'Hello john #2')
        self.assertEqual(next(self.results), 'Hello john #3')
        self.assertEqual(next(self.results), 'Hello john #4')
        self.assertEqual(self.results._limit, 15)
        # Second page.
        self.assertEqual(next(self.results), 'Hello john #5')
        self.assertEqual(next(self.results), 'Hello john #6')
        self.assertEqual(next(self.results), 'Hello john #7')
        self.assertEqual(next(self.results), 'Hello john #8')
        self.assertEqual(next(self.results), 'Hello john #9')
        self.assertEqual(self.results._limit, 10)
        # Third page.
        self.assertEqual(next(self.results), 'Hello john #10')
        self.assertEqual(next(self.results), 'Hello john #11')
        self.assertEqual(next(self.results), 'Hello john #12')
        self.assertRaises(StopIteration, self.results.next)
        self.assertEqual(self.results._limit, 7)

    def test_limit_smaller_than_first_page(self):
        results = ResultSet()
        results.to_call(fake_results, 'john', greeting='Hello', limit=2)
        self.assertEqual(next(results), 'Hello john #0')
        self.assertEqual(next(results), 'Hello john #1')
        self.assertRaises(StopIteration, results.next)

    def test_limit_equals_page(self):
        results = ResultSet()
        results.to_call(fake_results, 'john', greeting='Hello', limit=5)
        # First page
        self.assertEqual(next(results), 'Hello john #0')
        self.assertEqual(next(results), 'Hello john #1')
        self.assertEqual(next(results), 'Hello john #2')
        self.assertEqual(next(results), 'Hello john #3')
        self.assertEqual(next(results), 'Hello john #4')
        self.assertRaises(StopIteration, results.next)

    def test_limit_greater_than_page(self):
        results = ResultSet()
        results.to_call(fake_results, 'john', greeting='Hello', limit=6)
        # First page
        self.assertEqual(next(results), 'Hello john #0')
        self.assertEqual(next(results), 'Hello john #1')
        self.assertEqual(next(results), 'Hello john #2')
        self.assertEqual(next(results), 'Hello john #3')
        self.assertEqual(next(results), 'Hello john #4')
        # Second page
        self.assertEqual(next(results), 'Hello john #5')
        self.assertRaises(StopIteration, results.next)

    def test_iteration_noresults(self):
        def none(limit=10):
            return {
                'results': [],
            }

        results = ResultSet()
        results.to_call(none, limit=20)
        self.assertRaises(StopIteration, results.next)

    def test_iteration_sporadic_pages(self):
        # Some pages have no/incomplete results but have a ``LastEvaluatedKey``
        # (for instance, scans with filters), so we need to accommodate that.
        def sporadic():
            # A dict, because Python closures have read-only access to the
            # reference itself.
            count = {'value': -1}

            def _wrapper(limit=10, exclusive_start_key=None):
                count['value'] = count['value'] + 1

                if count['value'] == 0:
                    # Full page.
                    return {
                        'results': [
                            'Result #0',
                            'Result #1',
                            'Result #2',
                            'Result #3',
                        ],
                        'last_key': 'page-1'
                    }
                elif count['value'] == 1:
                    # Empty page but continue.
                    return {
                        'results': [],
                        'last_key': 'page-2'
                    }
                elif count['value'] == 2:
                    # Final page.
                    return {
                        'results': [
                            'Result #4',
                            'Result #5',
                            'Result #6',
                        ],
                    }

            return _wrapper

        results = ResultSet()
        results.to_call(sporadic(), limit=20)
        # First page
        self.assertEqual(next(results), 'Result #0')
        self.assertEqual(next(results), 'Result #1')
        self.assertEqual(next(results), 'Result #2')
        self.assertEqual(next(results), 'Result #3')
        # Second page (misses!)
        # Moves on to the third page
        self.assertEqual(next(results), 'Result #4')
        self.assertEqual(next(results), 'Result #5')
        self.assertEqual(next(results), 'Result #6')
        self.assertRaises(StopIteration, results.next)

    def test_list(self):
        self.assertEqual(list(self.results), [
            'Hello john #0',
            'Hello john #1',
            'Hello john #2',
            'Hello john #3',
            'Hello john #4',
            'Hello john #5',
            'Hello john #6',
            'Hello john #7',
            'Hello john #8',
            'Hello john #9',
            'Hello john #10',
            'Hello john #11',
            'Hello john #12'
        ])


def fake_batch_results(keys):
    results = []
    simulate_unprocessed = True

    if len(keys) and keys[0] == 'johndoe':
        simulate_unprocessed = False

    for key in keys:
        if simulate_unprocessed and key == 'johndoe':
            continue

        results.append("hello %s" % key)

    retval = {
        'results': results,
        'last_key': None,
    }

    if simulate_unprocessed:
        retval['unprocessed_keys'] = ['johndoe']

    return retval


class BatchGetResultSetTestCase(unittest.TestCase):
    def setUp(self):
        super(BatchGetResultSetTestCase, self).setUp()
        self.results = BatchGetResultSet(keys=[
            'alice',
            'bob',
            'jane',
            'johndoe',
        ])
        self.results.to_call(fake_batch_results)

    def test_fetch_more(self):
        # First "page".
        self.results.fetch_more()
        self.assertEqual(self.results._results, [
            'hello alice',
            'hello bob',
            'hello jane',
        ])
        self.assertEqual(self.results._keys_left, ['johndoe'])

        # Second "page".
        self.results.fetch_more()
        self.assertEqual(self.results._results, [
            'hello johndoe',
        ])

        # Empty "page". Nothing new gets added
        self.results.fetch_more()
        self.assertEqual(self.results._results, [])

        # Make sure we won't check for results in the future.
        self.assertFalse(self.results._results_left)

    def test_fetch_more_empty(self):
        self.results.to_call(lambda keys: {'results': [], 'last_key': None})

        self.results.fetch_more()
        self.assertEqual(self.results._results, [])
        self.assertRaises(StopIteration, self.results.next)

    def test_iteration(self):
        # First page.
        self.assertEqual(next(self.results), 'hello alice')
        self.assertEqual(next(self.results), 'hello bob')
        self.assertEqual(next(self.results), 'hello jane')
        self.assertEqual(next(self.results), 'hello johndoe')
        self.assertRaises(StopIteration, self.results.next)


class TableTestCase(unittest.TestCase):
    def setUp(self):
        super(TableTestCase, self).setUp()
        self.users = Table('users', connection=FakeDynamoDBConnection())
        self.default_connection = DynamoDBConnection(
            aws_access_key_id='access_key',
            aws_secret_access_key='secret_key'
        )

    def test__introspect_schema(self):
        raw_schema_1 = [
            {
                "AttributeName": "username",
                "KeyType": "HASH"
            },
            {
                "AttributeName": "date_joined",
                "KeyType": "RANGE"
            }
        ]
        raw_attributes_1 = [
            {
                'AttributeName': 'username',
                'AttributeType': 'S'
            },
            {
                'AttributeName': 'date_joined',
                'AttributeType': 'S'
            },
        ]
        schema_1 = self.users._introspect_schema(raw_schema_1, raw_attributes_1)
        self.assertEqual(len(schema_1), 2)
        self.assertTrue(isinstance(schema_1[0], HashKey))
        self.assertEqual(schema_1[0].name, 'username')
        self.assertTrue(isinstance(schema_1[1], RangeKey))
        self.assertEqual(schema_1[1].name, 'date_joined')

        raw_schema_2 = [
            {
                "AttributeName": "username",
                "KeyType": "BTREE"
            },
        ]
        raw_attributes_2 = [
            {
                'AttributeName': 'username',
                'AttributeType': 'S'
            },
        ]
        self.assertRaises(
            exceptions.UnknownSchemaFieldError,
            self.users._introspect_schema,
            raw_schema_2,
            raw_attributes_2
        )

        # Test a complex schema & ensure the types come back correctly.
        raw_schema_3 = [
            {
                "AttributeName": "user_id",
                "KeyType": "HASH"
            },
            {
                "AttributeName": "junk",
                "KeyType": "RANGE"
            }
        ]
        raw_attributes_3 = [
            {
                'AttributeName': 'user_id',
                'AttributeType': 'N'
            },
            {
                'AttributeName': 'junk',
                'AttributeType': 'B'
            },
        ]
        schema_3 = self.users._introspect_schema(raw_schema_3, raw_attributes_3)
        self.assertEqual(len(schema_3), 2)
        self.assertTrue(isinstance(schema_3[0], HashKey))
        self.assertEqual(schema_3[0].name, 'user_id')
        self.assertEqual(schema_3[0].data_type, NUMBER)
        self.assertTrue(isinstance(schema_3[1], RangeKey))
        self.assertEqual(schema_3[1].name, 'junk')
        self.assertEqual(schema_3[1].data_type, BINARY)

    def test__introspect_indexes(self):
        raw_indexes_1 = [
            {
                "IndexName": "MostRecentlyJoinedIndex",
                "KeySchema": [
                    {
                        "AttributeName": "username",
                        "KeyType": "HASH"
                    },
                    {
                        "AttributeName": "date_joined",
                        "KeyType": "RANGE"
                    }
                ],
                "Projection": {
                    "ProjectionType": "KEYS_ONLY"
                }
            },
            {
                "IndexName": "EverybodyIndex",
                "KeySchema": [
                    {
                        "AttributeName": "username",
                        "KeyType": "HASH"
                    },
                ],
                "Projection": {
                    "ProjectionType": "ALL"
                }
            },
            {
                "IndexName": "GenderIndex",
                "KeySchema": [
                    {
                        "AttributeName": "username",
                        "KeyType": "HASH"
                    },
                    {
                        "AttributeName": "date_joined",
                        "KeyType": "RANGE"
                    }
                ],
                "Projection": {
                    "ProjectionType": "INCLUDE",
                    "NonKeyAttributes": [
                        'gender',
                    ]
                }
            }
        ]
        indexes_1 = self.users._introspect_indexes(raw_indexes_1)
        self.assertEqual(len(indexes_1), 3)
        self.assertTrue(isinstance(indexes_1[0], KeysOnlyIndex))
        self.assertEqual(indexes_1[0].name, 'MostRecentlyJoinedIndex')
        self.assertEqual(len(indexes_1[0].parts), 2)
        self.assertTrue(isinstance(indexes_1[1], AllIndex))
        self.assertEqual(indexes_1[1].name, 'EverybodyIndex')
        self.assertEqual(len(indexes_1[1].parts), 1)
        self.assertTrue(isinstance(indexes_1[2], IncludeIndex))
        self.assertEqual(indexes_1[2].name, 'GenderIndex')
        self.assertEqual(len(indexes_1[2].parts), 2)
        self.assertEqual(indexes_1[2].includes_fields, ['gender'])

        raw_indexes_2 = [
            {
                "IndexName": "MostRecentlyJoinedIndex",
                "KeySchema": [
                    {
                        "AttributeName": "username",
                        "KeyType": "HASH"
                    },
                    {
                        "AttributeName": "date_joined",
                        "KeyType": "RANGE"
                    }
                ],
                "Projection": {
                    "ProjectionType": "SOMETHING_CRAZY"
                }
            },
        ]
        self.assertRaises(
            exceptions.UnknownIndexFieldError,
            self.users._introspect_indexes,
            raw_indexes_2
        )

    def test_initialization(self):
        users = Table('users', connection=self.default_connection)
        self.assertEqual(users.table_name, 'users')
        self.assertTrue(isinstance(users.connection, DynamoDBConnection))
        self.assertEqual(users.throughput['read'], 5)
        self.assertEqual(users.throughput['write'], 5)
        self.assertEqual(users.schema, None)
        self.assertEqual(users.indexes, None)

        groups = Table('groups', connection=FakeDynamoDBConnection())
        self.assertEqual(groups.table_name, 'groups')
        self.assertTrue(hasattr(groups.connection, 'assert_called_once_with'))

    def test_create_simple(self):
        conn = FakeDynamoDBConnection()

        with mock.patch.object(conn, 'create_table', return_value={}) \
                as mock_create_table:
            retval = Table.create('users', schema=[
                HashKey('username'),
                RangeKey('date_joined', data_type=NUMBER)
            ], connection=conn)
            self.assertTrue(retval)

        self.assertTrue(mock_create_table.called)
        mock_create_table.assert_called_once_with(attribute_definitions=[
            {
                'AttributeName': 'username',
                'AttributeType': 'S'
            },
            {
                'AttributeName': 'date_joined',
                'AttributeType': 'N'
            }
        ],
        table_name='users',
        key_schema=[
            {
                'KeyType': 'HASH',
                'AttributeName': 'username'
            },
            {
                'KeyType': 'RANGE',
                'AttributeName': 'date_joined'
            }
        ],
        provisioned_throughput={
            'WriteCapacityUnits': 5,
            'ReadCapacityUnits': 5
        })

    def test_create_full(self):
        conn = FakeDynamoDBConnection()

        with mock.patch.object(conn, 'create_table', return_value={}) \
                as mock_create_table:
            retval = Table.create('users', schema=[
                HashKey('username'),
                RangeKey('date_joined', data_type=NUMBER)
            ], throughput={
                'read':20,
                'write': 10,
            }, indexes=[
                KeysOnlyIndex('FriendCountIndex', parts=[
                    RangeKey('friend_count')
                ]),
            ], global_indexes=[
                GlobalKeysOnlyIndex('FullFriendCountIndex', parts=[
                    RangeKey('friend_count')
                ], throughput={
                    'read': 10,
                    'write': 8,
                }),
            ], connection=conn)
            self.assertTrue(retval)

        self.assertTrue(mock_create_table.called)
        mock_create_table.assert_called_once_with(attribute_definitions=[
            {
                'AttributeName': 'username',
                'AttributeType': 'S'
            },
            {
                'AttributeName': 'date_joined',
                'AttributeType': 'N'
            },
            {
                'AttributeName': 'friend_count',
                'AttributeType': 'S'
            }
        ],
        key_schema=[
            {
                'KeyType': 'HASH',
                'AttributeName': 'username'
            },
            {
                'KeyType': 'RANGE',
                'AttributeName': 'date_joined'
            }
        ],
        table_name='users',
        provisioned_throughput={
            'WriteCapacityUnits': 10,
            'ReadCapacityUnits': 20
        },
        global_secondary_indexes=[
            {
                'KeySchema': [
                    {
                        'KeyType': 'RANGE',
                        'AttributeName': 'friend_count'
                    }
                ],
                'IndexName': 'FullFriendCountIndex',
                'Projection': {
                    'ProjectionType': 'KEYS_ONLY'
                },
                'ProvisionedThroughput': {
                    'WriteCapacityUnits': 8,
                    'ReadCapacityUnits': 10
                }
            }
        ],
        local_secondary_indexes=[
            {
                'KeySchema': [
                    {
                        'KeyType': 'RANGE',
                        'AttributeName': 'friend_count'
                    }
                ],
                'IndexName': 'FriendCountIndex',
                'Projection': {
                    'ProjectionType': 'KEYS_ONLY'
                }
            }
        ])

    def test_describe(self):
        expected = {
            "Table": {
                "AttributeDefinitions": [
                    {
                        "AttributeName": "username",
                        "AttributeType": "S"
                    }
                ],
                "ItemCount": 5,
                "KeySchema": [
                    {
                        "AttributeName": "username",
                        "KeyType": "HASH"
                    }
                ],
                "LocalSecondaryIndexes": [
                    {
                        "IndexName": "UsernameIndex",
                        "KeySchema": [
                            {
                                "AttributeName": "username",
                                "KeyType": "HASH"
                            }
                        ],
                        "Projection": {
                            "ProjectionType": "KEYS_ONLY"
                        }
                    }
                ],
                "ProvisionedThroughput": {
                    "ReadCapacityUnits": 20,
                    "WriteCapacityUnits": 6
                },
                "TableName": "Thread",
                "TableStatus": "ACTIVE"
            }
        }

        with mock.patch.object(
                self.users.connection,
                'describe_table',
                return_value=expected) as mock_describe:
            self.assertEqual(self.users.throughput['read'], 5)
            self.assertEqual(self.users.throughput['write'], 5)
            self.assertEqual(self.users.schema, None)
            self.assertEqual(self.users.indexes, None)

            self.users.describe()

            self.assertEqual(self.users.throughput['read'], 20)
            self.assertEqual(self.users.throughput['write'], 6)
            self.assertEqual(len(self.users.schema), 1)
            self.assertEqual(isinstance(self.users.schema[0], HashKey), 1)
            self.assertEqual(len(self.users.indexes), 1)…