1from tests.compat import mock, unittest
2from boto.dynamodb2 import exceptions
3from boto.dynamodb2.fields import (HashKey, RangeKey,
4                                   AllIndex, KeysOnlyIndex, IncludeIndex,
5                                   GlobalAllIndex, GlobalKeysOnlyIndex,
6                                   GlobalIncludeIndex)
7from boto.dynamodb2.items import Item
8from boto.dynamodb2.layer1 import DynamoDBConnection
9from boto.dynamodb2.results import ResultSet, BatchGetResultSet
10from boto.dynamodb2.table import Table
11from boto.dynamodb2.types import (STRING, NUMBER, BINARY,
12                                  FILTER_OPERATORS, QUERY_OPERATORS)
13from boto.exception import JSONResponseError
14from boto.compat import six, long_type
15
16
17FakeDynamoDBConnection = mock.create_autospec(DynamoDBConnection)
18
19
20class SchemaFieldsTestCase(unittest.TestCase):
21    def test_hash_key(self):
22        hash_key = HashKey('hello')
23        self.assertEqual(hash_key.name, 'hello')
24        self.assertEqual(hash_key.data_type, STRING)
25        self.assertEqual(hash_key.attr_type, 'HASH')
26
27        self.assertEqual(hash_key.definition(), {
28            'AttributeName': 'hello',
29            'AttributeType': 'S'
30        })
31        self.assertEqual(hash_key.schema(), {
32            'AttributeName': 'hello',
33            'KeyType': 'HASH'
34        })
35
36    def test_range_key(self):
37        range_key = RangeKey('hello')
38        self.assertEqual(range_key.name, 'hello')
39        self.assertEqual(range_key.data_type, STRING)
40        self.assertEqual(range_key.attr_type, 'RANGE')
41
42        self.assertEqual(range_key.definition(), {
43            'AttributeName': 'hello',
44            'AttributeType': 'S'
45        })
46        self.assertEqual(range_key.schema(), {
47            'AttributeName': 'hello',
48            'KeyType': 'RANGE'
49        })
50
51    def test_alternate_type(self):
52        alt_key = HashKey('alt', data_type=NUMBER)
53        self.assertEqual(alt_key.name, 'alt')
54        self.assertEqual(alt_key.data_type, NUMBER)
55        self.assertEqual(alt_key.attr_type, 'HASH')
56
57        self.assertEqual(alt_key.definition(), {
58            'AttributeName': 'alt',
59            'AttributeType': 'N'
60        })
61        self.assertEqual(alt_key.schema(), {
62            'AttributeName': 'alt',
63            'KeyType': 'HASH'
64        })
65
66
67class IndexFieldTestCase(unittest.TestCase):
68    def test_all_index(self):
69        all_index = AllIndex('AllKeys', parts=[
70            HashKey('username'),
71            RangeKey('date_joined')
72        ])
73        self.assertEqual(all_index.name, 'AllKeys')
74        self.assertEqual([part.attr_type for part in all_index.parts], [
75            'HASH',
76            'RANGE'
77        ])
78        self.assertEqual(all_index.projection_type, 'ALL')
79
80        self.assertEqual(all_index.definition(), [
81            {'AttributeName': 'username', 'AttributeType': 'S'},
82            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
83        ])
84        self.assertEqual(all_index.schema(), {
85            'IndexName': 'AllKeys',
86            'KeySchema': [
87                {
88                    'AttributeName': 'username',
89                    'KeyType': 'HASH'
90                },
91                {
92                    'AttributeName': 'date_joined',
93                    'KeyType': 'RANGE'
94                }
95            ],
96            'Projection': {
97                'ProjectionType': 'ALL'
98            }
99        })
100
101    def test_keys_only_index(self):
102        keys_only = KeysOnlyIndex('KeysOnly', parts=[
103            HashKey('username'),
104            RangeKey('date_joined')
105        ])
106        self.assertEqual(keys_only.name, 'KeysOnly')
107        self.assertEqual([part.attr_type for part in keys_only.parts], [
108            'HASH',
109            'RANGE'
110        ])
111        self.assertEqual(keys_only.projection_type, 'KEYS_ONLY')
112
113        self.assertEqual(keys_only.definition(), [
114            {'AttributeName': 'username', 'AttributeType': 'S'},
115            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
116        ])
117        self.assertEqual(keys_only.schema(), {
118            'IndexName': 'KeysOnly',
119            'KeySchema': [
120                {
121                    'AttributeName': 'username',
122                    'KeyType': 'HASH'
123                },
124                {
125                    'AttributeName': 'date_joined',
126                    'KeyType': 'RANGE'
127                }
128            ],
129            'Projection': {
130                'ProjectionType': 'KEYS_ONLY'
131            }
132        })
133
134    def test_include_index(self):
135        include_index = IncludeIndex('IncludeKeys', parts=[
136            HashKey('username'),
137            RangeKey('date_joined')
138        ], includes=[
139            'gender',
140            'friend_count'
141        ])
142        self.assertEqual(include_index.name, 'IncludeKeys')
143        self.assertEqual([part.attr_type for part in include_index.parts], [
144            'HASH',
145            'RANGE'
146        ])
147        self.assertEqual(include_index.projection_type, 'INCLUDE')
148
149        self.assertEqual(include_index.definition(), [
150            {'AttributeName': 'username', 'AttributeType': 'S'},
151            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
152        ])
153        self.assertEqual(include_index.schema(), {
154            'IndexName': 'IncludeKeys',
155            'KeySchema': [
156                {
157                    'AttributeName': 'username',
158                    'KeyType': 'HASH'
159                },
160                {
161                    'AttributeName': 'date_joined',
162                    'KeyType': 'RANGE'
163                }
164            ],
165            'Projection': {
166                'ProjectionType': 'INCLUDE',
167                'NonKeyAttributes': [
168                    'gender',
169                    'friend_count',
170                ]
171            }
172        })
173
174    def test_global_all_index(self):
175        all_index = GlobalAllIndex('AllKeys', parts=[
176            HashKey('username'),
177            RangeKey('date_joined')
178        ],
179        throughput={
180            'read': 6,
181            'write': 2,
182        })
183        self.assertEqual(all_index.name, 'AllKeys')
184        self.assertEqual([part.attr_type for part in all_index.parts], [
185            'HASH',
186            'RANGE'
187        ])
188        self.assertEqual(all_index.projection_type, 'ALL')
189
190        self.assertEqual(all_index.definition(), [
191            {'AttributeName': 'username', 'AttributeType': 'S'},
192            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
193        ])
194        self.assertEqual(all_index.schema(), {
195            'IndexName': 'AllKeys',
196            'KeySchema': [
197                {
198                    'AttributeName': 'username',
199                    'KeyType': 'HASH'
200                },
201                {
202                    'AttributeName': 'date_joined',
203                    'KeyType': 'RANGE'
204                }
205            ],
206            'Projection': {
207                'ProjectionType': 'ALL'
208            },
209            'ProvisionedThroughput': {
210                'ReadCapacityUnits': 6,
211                'WriteCapacityUnits': 2
212            }
213        })
214
215    def test_global_keys_only_index(self):
216        keys_only = GlobalKeysOnlyIndex('KeysOnly', parts=[
217            HashKey('username'),
218            RangeKey('date_joined')
219        ],
220        throughput={
221            'read': 3,
222            'write': 4,
223        })
224        self.assertEqual(keys_only.name, 'KeysOnly')
225        self.assertEqual([part.attr_type for part in keys_only.parts], [
226            'HASH',
227            'RANGE'
228        ])
229        self.assertEqual(keys_only.projection_type, 'KEYS_ONLY')
230
231        self.assertEqual(keys_only.definition(), [
232            {'AttributeName': 'username', 'AttributeType': 'S'},
233            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
234        ])
235        self.assertEqual(keys_only.schema(), {
236            'IndexName': 'KeysOnly',
237            'KeySchema': [
238                {
239                    'AttributeName': 'username',
240                    'KeyType': 'HASH'
241                },
242                {
243                    'AttributeName': 'date_joined',
244                    'KeyType': 'RANGE'
245                }
246            ],
247            'Projection': {
248                'ProjectionType': 'KEYS_ONLY'
249            },
250            'ProvisionedThroughput': {
251                'ReadCapacityUnits': 3,
252                'WriteCapacityUnits': 4
253            }
254        })
255
256    def test_global_include_index(self):
257        # Lean on the default throughput
258        include_index = GlobalIncludeIndex('IncludeKeys', parts=[
259            HashKey('username'),
260            RangeKey('date_joined')
261        ], includes=[
262            'gender',
263            'friend_count'
264        ])
265        self.assertEqual(include_index.name, 'IncludeKeys')
266        self.assertEqual([part.attr_type for part in include_index.parts], [
267            'HASH',
268            'RANGE'
269        ])
270        self.assertEqual(include_index.projection_type, 'INCLUDE')
271
272        self.assertEqual(include_index.definition(), [
273            {'AttributeName': 'username', 'AttributeType': 'S'},
274            {'AttributeName': 'date_joined', 'AttributeType': 'S'}
275        ])
276        self.assertEqual(include_index.schema(), {
277            'IndexName': 'IncludeKeys',
278            'KeySchema': [
279                {
280                    'AttributeName': 'username',
281                    'KeyType': 'HASH'
282                },
283                {
284                    'AttributeName': 'date_joined',
285                    'KeyType': 'RANGE'
286                }
287            ],
288            'Projection': {
289                'ProjectionType': 'INCLUDE',
290                'NonKeyAttributes': [
291                    'gender',
292                    'friend_count',
293                ]
294            },
295            'ProvisionedThroughput': {
296                'ReadCapacityUnits': 5,
297                'WriteCapacityUnits': 5
298            }
299        })
300
301    def test_global_include_index_throughput(self):
302        include_index = GlobalIncludeIndex('IncludeKeys', parts=[
303            HashKey('username'),
304            RangeKey('date_joined')
305        ], includes=[
306            'gender',
307            'friend_count'
308        ], throughput={
309            'read': 10,
310            'write': 8
311        })
312
313        self.assertEqual(include_index.schema(), {
314            'IndexName': 'IncludeKeys',
315            'KeySchema': [
316                {
317                    'AttributeName': 'username',
318                    'KeyType': 'HASH'
319                },
320                {
321                    'AttributeName': 'date_joined',
322                    'KeyType': 'RANGE'
323                }
324            ],
325            'Projection': {
326                'ProjectionType': 'INCLUDE',
327                'NonKeyAttributes': [
328                    'gender',
329                    'friend_count',
330                ]
331            },
332            'ProvisionedThroughput': {
333                'ReadCapacityUnits': 10,
334                'WriteCapacityUnits': 8
335            }
336        })
337
338
339class ItemTestCase(unittest.TestCase):
340    if six.PY2:
341        assertCountEqual = unittest.TestCase.assertItemsEqual
342
343    def setUp(self):
344        super(ItemTestCase, self).setUp()
345        self.table = Table('whatever', connection=FakeDynamoDBConnection())
346        self.johndoe = self.create_item({
347            'username': 'johndoe',
348            'first_name': 'John',
349            'date_joined': 12345,
350        })
351
352    def create_item(self, data):
353        return Item(self.table, data=data)
354
355    def test_initialization(self):
356        empty_item = Item(self.table)
357        self.assertEqual(empty_item.table, self.table)
358        self.assertEqual(empty_item._data, {})
359
360        full_item = Item(self.table, data={
361            'username': 'johndoe',
362            'date_joined': 12345,
363        })
364        self.assertEqual(full_item.table, self.table)
365        self.assertEqual(full_item._data, {
366            'username': 'johndoe',
367            'date_joined': 12345,
368        })
369
370    # The next couple methods make use of ``sorted(...)`` so we get consistent
371    # ordering everywhere & no erroneous failures.
372
373    def test_keys(self):
374        self.assertCountEqual(self.johndoe.keys(), [
375            'date_joined',
376            'first_name',
377            'username',
378        ])
379
380    def test_values(self):
381        self.assertCountEqual(self.johndoe.values(),
382                              [12345, 'John', 'johndoe'])
383
384    def test_contains(self):
385        self.assertIn('username', self.johndoe)
386        self.assertIn('first_name', self.johndoe)
387        self.assertIn('date_joined', self.johndoe)
388        self.assertNotIn('whatever', self.johndoe)
389
390    def test_iter(self):
391        self.assertCountEqual(self.johndoe,
392                              ['johndoe', 'John', 12345])
393
394    def test_get(self):
395        self.assertEqual(self.johndoe.get('username'), 'johndoe')
396        self.assertEqual(self.johndoe.get('first_name'), 'John')
397        self.assertEqual(self.johndoe.get('date_joined'), 12345)
398
399        # Test a missing key. No default yields ``None``.
400        self.assertEqual(self.johndoe.get('last_name'), None)
401        # This time with a default.
402        self.assertEqual(self.johndoe.get('last_name', True), True)
403
404    def test_items(self):
405        self.assertCountEqual(
406            self.johndoe.items(),
407            [
408                ('date_joined', 12345),
409                ('first_name', 'John'),
410                ('username', 'johndoe'),
411            ])
412
413    def test_attribute_access(self):
414        self.assertEqual(self.johndoe['username'], 'johndoe')
415        self.assertEqual(self.johndoe['first_name'], 'John')
416        self.assertEqual(self.johndoe['date_joined'], 12345)
417
418        # Test a missing key.
419        self.assertEqual(self.johndoe['last_name'], None)
420
421        # Set a key.
422        self.johndoe['last_name'] = 'Doe'
423        # Test accessing the new key.
424        self.assertEqual(self.johndoe['last_name'], 'Doe')
425
426        # Delete a key.
427        del self.johndoe['last_name']
428        # Test the now-missing-again key.
429        self.assertEqual(self.johndoe['last_name'], None)
430
431    def test_needs_save(self):
432        self.johndoe.mark_clean()
433        self.assertFalse(self.johndoe.needs_save())
434        self.johndoe['last_name'] = 'Doe'
435        self.assertTrue(self.johndoe.needs_save())
436
437    def test_needs_save_set_changed(self):
438        # First, ensure we're clean.
439        self.johndoe.mark_clean()
440        self.assertFalse(self.johndoe.needs_save())
441        # Add a friends collection.
442        self.johndoe['friends'] = set(['jane', 'alice'])
443        self.assertTrue(self.johndoe.needs_save())
444        # Now mark it clean, then change the collection.
445        # This does NOT call ``__setitem__``, so the item used to be
446        # incorrectly appearing to be clean, when it had in fact been changed.
447        self.johndoe.mark_clean()
448        self.assertFalse(self.johndoe.needs_save())
449        self.johndoe['friends'].add('bob')
450        self.assertTrue(self.johndoe.needs_save())
451
452    def test_mark_clean(self):
453        self.johndoe['last_name'] = 'Doe'
454        self.assertTrue(self.johndoe.needs_save())
455        self.johndoe.mark_clean()
456        self.assertFalse(self.johndoe.needs_save())
457
458    def test_load(self):
459        empty_item = Item(self.table)
460        empty_item.load({
461            'Item': {
462                'username': {'S': 'johndoe'},
463                'first_name': {'S': 'John'},
464                'last_name': {'S': 'Doe'},
465                'date_joined': {'N': '1366056668'},
466                'friend_count': {'N': '3'},
467                'friends': {'SS': ['alice', 'bob', 'jane']},
468            }
469        })
470        self.assertEqual(empty_item['username'], 'johndoe')
471        self.assertEqual(empty_item['date_joined'], 1366056668)
472        self.assertEqual(sorted(empty_item['friends']), sorted([
473            'alice',
474            'bob',
475            'jane'
476        ]))
477
478    def test_get_keys(self):
479        # Setup the data.
480        self.table.schema = [
481            HashKey('username'),
482            RangeKey('date_joined'),
483        ]
484        self.assertEqual(self.johndoe.get_keys(), {
485            'username': 'johndoe',
486            'date_joined': 12345,
487        })
488
489    def test_get_raw_keys(self):
490        # Setup the data.
491        self.table.schema = [
492            HashKey('username'),
493            RangeKey('date_joined'),
494        ]
495        self.assertEqual(self.johndoe.get_raw_keys(), {
496            'username': {'S': 'johndoe'},
497            'date_joined': {'N': '12345'},
498        })
499
500    def test_build_expects(self):
501        # Pristine.
502        self.assertEqual(self.johndoe.build_expects(), {
503            'first_name': {
504                'Exists': False,
505            },
506            'username': {
507                'Exists': False,
508            },
509            'date_joined': {
510                'Exists': False,
511            },
512        })
513
514        # Without modifications.
515        self.johndoe.mark_clean()
516        self.assertEqual(self.johndoe.build_expects(), {
517            'first_name': {
518                'Exists': True,
519                'Value': {
520                    'S': 'John',
521                },
522            },
523            'username': {
524                'Exists': True,
525                'Value': {
526                    'S': 'johndoe',
527                },
528            },
529            'date_joined': {
530                'Exists': True,
531                'Value': {
532                    'N': '12345',
533                },
534            },
535        })
536
537        # Change some data.
538        self.johndoe['first_name'] = 'Johann'
539        # Add some data.
540        self.johndoe['last_name'] = 'Doe'
541        # Delete some data.
542        del self.johndoe['date_joined']
543
544        # All fields (default).
545        self.assertEqual(self.johndoe.build_expects(), {
546            'first_name': {
547                'Exists': True,
548                'Value': {
549                    'S': 'John',
550                },
551            },
552            'last_name': {
553                'Exists': False,
554            },
555            'username': {
556                'Exists': True,
557                'Value': {
558                    'S': 'johndoe',
559                },
560            },
561            'date_joined': {
562                'Exists': True,
563                'Value': {
564                    'N': '12345',
565                },
566            },
567        })
568
569        # Only a subset of the fields.
570        self.assertEqual(self.johndoe.build_expects(fields=[
571            'first_name',
572            'last_name',
573            'date_joined',
574        ]), {
575            'first_name': {
576                'Exists': True,
577                'Value': {
578                    'S': 'John',
579                },
580            },
581            'last_name': {
582                'Exists': False,
583            },
584            'date_joined': {
585                'Exists': True,
586                'Value': {
587                    'N': '12345',
588                },
589            },
590        })
591
592    def test_prepare_full(self):
593        self.assertEqual(self.johndoe.prepare_full(), {
594            'username': {'S': 'johndoe'},
595            'first_name': {'S': 'John'},
596            'date_joined': {'N': '12345'}
597        })
598
599        self.johndoe['friends'] = set(['jane', 'alice'])
600        data = self.johndoe.prepare_full()
601        self.assertEqual(data['username'], {'S': 'johndoe'})
602        self.assertEqual(data['first_name'], {'S': 'John'})
603        self.assertEqual(data['date_joined'], {'N': '12345'})
604        self.assertCountEqual(data['friends']['SS'],
605                              ['jane', 'alice'])
606
607    def test_prepare_full_empty_set(self):
608        self.johndoe['friends'] = set()
609        self.assertEqual(self.johndoe.prepare_full(), {
610            'username': {'S': 'johndoe'},
611            'first_name': {'S': 'John'},
612            'date_joined': {'N': '12345'}
613        })
614
615    def test_prepare_partial(self):
616        self.johndoe.mark_clean()
617        # Change some data.
618        self.johndoe['first_name'] = 'Johann'
619        # Add some data.
620        self.johndoe['last_name'] = 'Doe'
621        # Delete some data.
622        del self.johndoe['date_joined']
623
624        final_data, fields = self.johndoe.prepare_partial()
625        self.assertEqual(final_data, {
626            'date_joined': {
627                'Action': 'DELETE',
628            },
629            'first_name': {
630                'Action': 'PUT',
631                'Value': {'S': 'Johann'},
632            },
633            'last_name': {
634                'Action': 'PUT',
635                'Value': {'S': 'Doe'},
636            },
637        })
638        self.assertEqual(fields, set([
639            'first_name',
640            'last_name',
641            'date_joined'
642        ]))
643
644    def test_prepare_partial_empty_set(self):
645        self.johndoe.mark_clean()
646        # Change some data.
647        self.johndoe['first_name'] = 'Johann'
648        # Add some data.
649        self.johndoe['last_name'] = 'Doe'
650        # Delete some data.
651        del self.johndoe['date_joined']
652        # Put an empty set on the ``Item``.
653        self.johndoe['friends'] = set()
654
655        final_data, fields = self.johndoe.prepare_partial()
656        self.assertEqual(final_data, {
657            'date_joined': {
658                'Action': 'DELETE',
659            },
660            'first_name': {
661                'Action': 'PUT',
662                'Value': {'S': 'Johann'},
663            },
664            'last_name': {
665                'Action': 'PUT',
666                'Value': {'S': 'Doe'},
667            },
668        })
669        self.assertEqual(fields, set([
670            'first_name',
671            'last_name',
672            'date_joined'
673        ]))
674
675    def test_save_no_changes(self):
676        # Unchanged, no save.
677        with mock.patch.object(self.table, '_put_item', return_value=True) \
678                as mock_put_item:
679            # Pretend we loaded it via ``get_item``...
680            self.johndoe.mark_clean()
681            self.assertFalse(self.johndoe.save())
682
683        self.assertFalse(mock_put_item.called)
684
685    def test_save_with_changes(self):
686        # With changed data.
687        with mock.patch.object(self.table, '_put_item', return_value=True) \
688                as mock_put_item:
689            self.johndoe.mark_clean()
690            self.johndoe['first_name'] = 'J'
691            self.johndoe['new_attr'] = 'never_seen_before'
692            self.assertTrue(self.johndoe.save())
693            self.assertFalse(self.johndoe.needs_save())
694
695        self.assertTrue(mock_put_item.called)
696        mock_put_item.assert_called_once_with({
697            'username': {'S': 'johndoe'},
698            'first_name': {'S': 'J'},
699            'new_attr': {'S': 'never_seen_before'},
700            'date_joined': {'N': '12345'}
701        }, expects={
702            'username': {
703                'Value': {
704                    'S': 'johndoe',
705                },
706                'Exists': True,
707            },
708            'first_name': {
709                'Value': {
710                    'S': 'John',
711                },
712                'Exists': True,
713            },
714            'new_attr': {
715                'Exists': False,
716            },
717            'date_joined': {
718                'Value': {
719                    'N': '12345',
720                },
721                'Exists': True,
722            },
723        })
724
725    def test_save_with_changes_overwrite(self):
726        # With changed data.
727        with mock.patch.object(self.table, '_put_item', return_value=True) \
728                as mock_put_item:
729            self.johndoe['first_name'] = 'J'
730            self.johndoe['new_attr'] = 'never_seen_before'
731            # OVERWRITE ALL THE THINGS
732            self.assertTrue(self.johndoe.save(overwrite=True))
733            self.assertFalse(self.johndoe.needs_save())
734
735        self.assertTrue(mock_put_item.called)
736        mock_put_item.assert_called_once_with({
737            'username': {'S': 'johndoe'},
738            'first_name': {'S': 'J'},
739            'new_attr': {'S': 'never_seen_before'},
740            'date_joined': {'N': '12345'}
741        }, expects=None)
742
743    def test_partial_no_changes(self):
744        # Unchanged, no save.
745        with mock.patch.object(self.table, '_update_item', return_value=True) \
746                as mock_update_item:
747            # Pretend we loaded it via ``get_item``...
748            self.johndoe.mark_clean()
749            self.assertFalse(self.johndoe.partial_save())
750
751        self.assertFalse(mock_update_item.called)
752
753    def test_partial_with_changes(self):
754        # Setup the data.
755        self.table.schema = [
756            HashKey('username'),
757        ]
758
759        # With changed data.
760        with mock.patch.object(self.table, '_update_item', return_value=True) \
761                as mock_update_item:
762            # Pretend we loaded it via ``get_item``...
763            self.johndoe.mark_clean()
764            # Now... MODIFY!!!
765            self.johndoe['first_name'] = 'J'
766            self.johndoe['last_name'] = 'Doe'
767            del self.johndoe['date_joined']
768            self.assertTrue(self.johndoe.partial_save())
769            self.assertFalse(self.johndoe.needs_save())
770
771        self.assertTrue(mock_update_item.called)
772        mock_update_item.assert_called_once_with({
773            'username': 'johndoe',
774        }, {
775            'first_name': {
776                'Action': 'PUT',
777                'Value': {'S': 'J'},
778            },
779            'last_name': {
780                'Action': 'PUT',
781                'Value': {'S': 'Doe'},
782            },
783            'date_joined': {
784                'Action': 'DELETE',
785            }
786        }, expects={
787            'first_name': {
788                'Value': {
789                    'S': 'John',
790                },
791                'Exists': True
792            },
793            'last_name': {
794                'Exists': False
795            },
796            'date_joined': {
797                'Value': {
798                    'N': '12345',
799                },
800                'Exists': True
801            },
802        })
803
804    def test_delete(self):
805        # Setup the data.
806        self.table.schema = [
807            HashKey('username'),
808            RangeKey('date_joined'),
809        ]
810
811        with mock.patch.object(self.table, 'delete_item', return_value=True) \
812                as mock_delete_item:
813            self.johndoe.delete()
814
815        self.assertTrue(mock_delete_item.called)
816        mock_delete_item.assert_called_once_with(
817            username='johndoe',
818            date_joined=12345
819        )
820
821    def test_nonzero(self):
822        self.assertTrue(self.johndoe)
823        self.assertFalse(self.create_item({}))
824
825
826class ItemFromItemTestCase(ItemTestCase):
827    def setUp(self):
828        super(ItemFromItemTestCase, self).setUp()
829        self.johndoe = self.create_item(self.johndoe)
830
831
832def fake_results(name, greeting='hello', exclusive_start_key=None, limit=None):
833    if exclusive_start_key is None:
834        exclusive_start_key = -1
835
836    if limit == 0:
837        raise Exception("Web Service Returns '400 Bad Request'")
838
839    end_cap = 13
840    results = []
841    start_key = exclusive_start_key + 1
842
843    for i in range(start_key, start_key + 5):
844        if i < end_cap:
845            results.append("%s %s #%s" % (greeting, name, i))
846
847    # Don't return more than limit results
848    if limit < len(results):
849        results = results[:limit]
850
851    retval = {
852        'results': results,
853    }
854
855    if exclusive_start_key + 5 < end_cap:
856        retval['last_key'] = exclusive_start_key + 5
857
858    return retval
859
860
861class ResultSetTestCase(unittest.TestCase):
862    def setUp(self):
863        super(ResultSetTestCase, self).setUp()
864        self.results = ResultSet()
865        self.result_function = mock.MagicMock(side_effect=fake_results)
866        self.results.to_call(self.result_function, 'john', greeting='Hello', limit=20)
867
868    def test_first_key(self):
869        self.assertEqual(self.results.first_key, 'exclusive_start_key')
870
871    def test_max_page_size_fetch_more(self):
872        self.results = ResultSet(max_page_size=10)
873        self.results.to_call(self.result_function, 'john', greeting='Hello')
874        self.results.fetch_more()
875        self.result_function.assert_called_with('john', greeting='Hello', limit=10)
876        self.result_function.reset_mock()
877
878    def test_max_page_size_and_smaller_limit_fetch_more(self):
879        self.results = ResultSet(max_page_size=10)
880        self.results.to_call(self.result_function, 'john', greeting='Hello', limit=5)
881        self.results.fetch_more()
882        self.result_function.assert_called_with('john', greeting='Hello', limit=5)
883        self.result_function.reset_mock()
884
885    def test_max_page_size_and_bigger_limit_fetch_more(self):
886        self.results = ResultSet(max_page_size=10)
887        self.results.to_call(self.result_function, 'john', greeting='Hello', limit=15)
888        self.results.fetch_more()
889        self.result_function.assert_called_with('john', greeting='Hello', limit=10)
890        self.result_function.reset_mock()
891
892    def test_fetch_more(self):
893        # First "page".
894        self.results.fetch_more()
895        self.assertEqual(self.results._results, [
896            'Hello john #0',
897            'Hello john #1',
898            'Hello john #2',
899            'Hello john #3',
900            'Hello john #4',
901        ])
902
903        self.result_function.assert_called_with('john', greeting='Hello', limit=20)
904        self.result_function.reset_mock()
905
906        # Fake in a last key.
907        self.results._last_key_seen = 4
908        # Second "page".
909        self.results.fetch_more()
910        self.assertEqual(self.results._results, [
911            'Hello john #5',
912            'Hello john #6',
913            'Hello john #7',
914            'Hello john #8',
915            'Hello john #9',
916        ])
917
918        self.result_function.assert_called_with('john', greeting='Hello', limit=20, exclusive_start_key=4)
919        self.result_function.reset_mock()
920
921        # Fake in a last key.
922        self.results._last_key_seen = 9
923        # Last "page".
924        self.results.fetch_more()
925        self.assertEqual(self.results._results, [
926            'Hello john #10',
927            'Hello john #11',
928            'Hello john #12',
929        ])
930
931        # Fake in a key outside the range.
932        self.results._last_key_seen = 15
933        # Empty "page". Nothing new gets added
934        self.results.fetch_more()
935        self.assertEqual(self.results._results, [])
936
937        # Make sure we won't check for results in the future.
938        self.assertFalse(self.results._results_left)
939
940    def test_iteration(self):
941        # First page.
942        self.assertEqual(next(self.results), 'Hello john #0')
943        self.assertEqual(next(self.results), 'Hello john #1')
944        self.assertEqual(next(self.results), 'Hello john #2')
945        self.assertEqual(next(self.results), 'Hello john #3')
946        self.assertEqual(next(self.results), 'Hello john #4')
947        self.assertEqual(self.results._limit, 15)
948        # Second page.
949        self.assertEqual(next(self.results), 'Hello john #5')
950        self.assertEqual(next(self.results), 'Hello john #6')
951        self.assertEqual(next(self.results), 'Hello john #7')
952        self.assertEqual(next(self.results), 'Hello john #8')
953        self.assertEqual(next(self.results), 'Hello john #9')
954        self.assertEqual(self.results._limit, 10)
955        # Third page.
956        self.assertEqual(next(self.results), 'Hello john #10')
957        self.assertEqual(next(self.results), 'Hello john #11')
958        self.assertEqual(next(self.results), 'Hello john #12')
959        self.assertRaises(StopIteration, self.results.next)
960        self.assertEqual(self.results._limit, 7)
961
962    def test_limit_smaller_than_first_page(self):
963        results = ResultSet()
964        results.to_call(fake_results, 'john', greeting='Hello', limit=2)
965        self.assertEqual(next(results), 'Hello john #0')
966        self.assertEqual(next(results), 'Hello john #1')
967        self.assertRaises(StopIteration, results.next)
968
969    def test_limit_equals_page(self):
970        results = ResultSet()
971        results.to_call(fake_results, 'john', greeting='Hello', limit=5)
972        # First page
973        self.assertEqual(next(results), 'Hello john #0')
974        self.assertEqual(next(results), 'Hello john #1')
975        self.assertEqual(next(results), 'Hello john #2')
976        self.assertEqual(next(results), 'Hello john #3')
977        self.assertEqual(next(results), 'Hello john #4')
978        self.assertRaises(StopIteration, results.next)
979
980    def test_limit_greater_than_page(self):
981        results = ResultSet()
982        results.to_call(fake_results, 'john', greeting='Hello', limit=6)
983        # First page
984        self.assertEqual(next(results), 'Hello john #0')
985        self.assertEqual(next(results), 'Hello john #1')
986        self.assertEqual(next(results), 'Hello john #2')
987        self.assertEqual(next(results), 'Hello john #3')
988        self.assertEqual(next(results), 'Hello john #4')
989        # Second page
990        self.assertEqual(next(results), 'Hello john #5')
991        self.assertRaises(StopIteration, results.next)
992
993    def test_iteration_noresults(self):
994        def none(limit=10):
995            return {
996                'results': [],
997            }
998
999        results = ResultSet()
1000        results.to_call(none, limit=20)
1001        self.assertRaises(StopIteration, results.next)
1002
1003    def test_iteration_sporadic_pages(self):
1004        # Some pages have no/incomplete results but have a ``LastEvaluatedKey``
1005        # (for instance, scans with filters), so we need to accommodate that.
1006        def sporadic():
1007            # A dict, because Python closures have read-only access to the
1008            # reference itself.
1009            count = {'value': -1}
1010
1011            def _wrapper(limit=10, exclusive_start_key=None):
1012                count['value'] = count['value'] + 1
1013
1014                if count['value'] == 0:
1015                    # Full page.
1016                    return {
1017                        'results': [
1018                            'Result #0',
1019                            'Result #1',
1020                            'Result #2',
1021                            'Result #3',
1022                        ],
1023                        'last_key': 'page-1'
1024                    }
1025                elif count['value'] == 1:
1026                    # Empty page but continue.
1027                    return {
1028                        'results': [],
1029                        'last_key': 'page-2'
1030                    }
1031                elif count['value'] == 2:
1032                    # Final page.
1033                    return {
1034                        'results': [
1035                            'Result #4',
1036                            'Result #5',
1037                            'Result #6',
1038                        ],
1039                    }
1040
1041            return _wrapper
1042
1043        results = ResultSet()
1044        results.to_call(sporadic(), limit=20)
1045        # First page
1046        self.assertEqual(next(results), 'Result #0')
1047        self.assertEqual(next(results), 'Result #1')
1048        self.assertEqual(next(results), 'Result #2')
1049        self.assertEqual(next(results), 'Result #3')
1050        # Second page (misses!)
1051        # Moves on to the third page
1052        self.assertEqual(next(results), 'Result #4')
1053        self.assertEqual(next(results), 'Result #5')
1054        self.assertEqual(next(results), 'Result #6')
1055        self.assertRaises(StopIteration, results.next)
1056
1057    def test_list(self):
1058        self.assertEqual(list(self.results), [
1059            'Hello john #0',
1060            'Hello john #1',
1061            'Hello john #2',
1062            'Hello john #3',
1063            'Hello john #4',
1064            'Hello john #5',
1065            'Hello john #6',
1066            'Hello john #7',
1067            'Hello john #8',
1068            'Hello john #9',
1069            'Hello john #10',
1070            'Hello john #11',
1071            'Hello john #12'
1072        ])
1073
1074
1075def fake_batch_results(keys):
1076    results = []
1077    simulate_unprocessed = True
1078
1079    if len(keys) and keys[0] == 'johndoe':
1080        simulate_unprocessed = False
1081
1082    for key in keys:
1083        if simulate_unprocessed and key == 'johndoe':
1084            continue
1085
1086        results.append("hello %s" % key)
1087
1088    retval = {
1089        'results': results,
1090        'last_key': None,
1091    }
1092
1093    if simulate_unprocessed:
1094        retval['unprocessed_keys'] = ['johndoe']
1095
1096    return retval
1097
1098
1099class BatchGetResultSetTestCase(unittest.TestCase):
1100    def setUp(self):
1101        super(BatchGetResultSetTestCase, self).setUp()
1102        self.results = BatchGetResultSet(keys=[
1103            'alice',
1104            'bob',
1105            'jane',
1106            'johndoe',
1107        ])
1108        self.results.to_call(fake_batch_results)
1109
1110    def test_fetch_more(self):
1111        # First "page".
1112        self.results.fetch_more()
1113        self.assertEqual(self.results._results, [
1114            'hello alice',
1115            'hello bob',
1116            'hello jane',
1117        ])
1118        self.assertEqual(self.results._keys_left, ['johndoe'])
1119
1120        # Second "page".
1121        self.results.fetch_more()
1122        self.assertEqual(self.results._results, [
1123            'hello johndoe',
1124        ])
1125
1126        # Empty "page". Nothing new gets added
1127        self.results.fetch_more()
1128        self.assertEqual(self.results._results, [])
1129
1130        # Make sure we won't check for results in the future.
1131        self.assertFalse(self.results._results_left)
1132
1133    def test_fetch_more_empty(self):
1134        self.results.to_call(lambda keys: {'results': [], 'last_key': None})
1135
1136        self.results.fetch_more()
1137        self.assertEqual(self.results._results, [])
1138        self.assertRaises(StopIteration, self.results.next)
1139
1140    def test_iteration(self):
1141        # First page.
1142        self.assertEqual(next(self.results), 'hello alice')
1143        self.assertEqual(next(self.results), 'hello bob')
1144        self.assertEqual(next(self.results), 'hello jane')
1145        self.assertEqual(next(self.results), 'hello johndoe')
1146        self.assertRaises(StopIteration, self.results.next)
1147
1148
1149class TableTestCase(unittest.TestCase):
1150    def setUp(self):
1151        super(TableTestCase, self).setUp()
1152        self.users = Table('users', connection=FakeDynamoDBConnection())
1153        self.default_connection = DynamoDBConnection(
1154            aws_access_key_id='access_key',
1155            aws_secret_access_key='secret_key'
1156        )
1157
1158    def test__introspect_schema(self):
1159        raw_schema_1 = [
1160            {
1161                "AttributeName": "username",
1162                "KeyType": "HASH"
1163            },
1164            {
1165                "AttributeName": "date_joined",
1166                "KeyType": "RANGE"
1167            }
1168        ]
1169        raw_attributes_1 = [
1170            {
1171                'AttributeName': 'username',
1172                'AttributeType': 'S'
1173            },
1174            {
1175                'AttributeName': 'date_joined',
1176                'AttributeType': 'S'
1177            },
1178        ]
1179        schema_1 = self.users._introspect_schema(raw_schema_1, raw_attributes_1)
1180        self.assertEqual(len(schema_1), 2)
1181        self.assertTrue(isinstance(schema_1[0], HashKey))
1182        self.assertEqual(schema_1[0].name, 'username')
1183        self.assertTrue(isinstance(schema_1[1], RangeKey))
1184        self.assertEqual(schema_1[1].name, 'date_joined')
1185
1186        raw_schema_2 = [
1187            {
1188                "AttributeName": "username",
1189                "KeyType": "BTREE"
1190            },
1191        ]
1192        raw_attributes_2 = [
1193            {
1194                'AttributeName': 'username',
1195                'AttributeType': 'S'
1196            },
1197        ]
1198        self.assertRaises(
1199            exceptions.UnknownSchemaFieldError,
1200            self.users._introspect_schema,
1201            raw_schema_2,
1202            raw_attributes_2
1203        )
1204
1205        # Test a complex schema & ensure the types come back correctly.
1206        raw_schema_3 = [
1207            {
1208                "AttributeName": "user_id",
1209                "KeyType": "HASH"
1210            },
1211            {
1212                "AttributeName": "junk",
1213                "KeyType": "RANGE"
1214            }
1215        ]
1216        raw_attributes_3 = [
1217            {
1218                'AttributeName': 'user_id',
1219                'AttributeType': 'N'
1220            },
1221            {
1222                'AttributeName': 'junk',
1223                'AttributeType': 'B'
1224            },
1225        ]
1226        schema_3 = self.users._introspect_schema(raw_schema_3, raw_attributes_3)
1227        self.assertEqual(len(schema_3), 2)
1228        self.assertTrue(isinstance(schema_3[0], HashKey))
1229        self.assertEqual(schema_3[0].name, 'user_id')
1230        self.assertEqual(schema_3[0].data_type, NUMBER)
1231        self.assertTrue(isinstance(schema_3[1], RangeKey))
1232        self.assertEqual(schema_3[1].name, 'junk')
1233        self.assertEqual(schema_3[1].data_type, BINARY)
1234
1235    def test__introspect_indexes(self):
1236        raw_indexes_1 = [
1237            {
1238                "IndexName": "MostRecentlyJoinedIndex",
1239                "KeySchema": [
1240                    {
1241                        "AttributeName": "username",
1242                        "KeyType": "HASH"
1243                    },
1244                    {
1245                        "AttributeName": "date_joined",
1246                        "KeyType": "RANGE"
1247                    }
1248                ],
1249                "Projection": {
1250                    "ProjectionType": "KEYS_ONLY"
1251                }
1252            },
1253            {
1254                "IndexName": "EverybodyIndex",
1255                "KeySchema": [
1256                    {
1257                        "AttributeName": "username",
1258                        "KeyType": "HASH"
1259                    },
1260                ],
1261                "Projection": {
1262                    "ProjectionType": "ALL"
1263                }
1264            },
1265            {
1266                "IndexName": "GenderIndex",
1267                "KeySchema": [
1268                    {
1269                        "AttributeName": "username",
1270                        "KeyType": "HASH"
1271                    },
1272                    {
1273                        "AttributeName": "date_joined",
1274                        "KeyType": "RANGE"
1275                    }
1276                ],
1277                "Projection": {
1278                    "ProjectionType": "INCLUDE",
1279                    "NonKeyAttributes": [
1280                        'gender',
1281                    ]
1282                }
1283            }
1284        ]
1285        indexes_1 = self.users._introspect_indexes(raw_indexes_1)
1286        self.assertEqual(len(indexes_1), 3)
1287        self.assertTrue(isinstance(indexes_1[0], KeysOnlyIndex))
1288        self.assertEqual(indexes_1[0].name, 'MostRecentlyJoinedIndex')
1289        self.assertEqual(len(indexes_1[0].parts), 2)
1290        self.assertTrue(isinstance(indexes_1[1], AllIndex))
1291        self.assertEqual(indexes_1[1].name, 'EverybodyIndex')
1292        self.assertEqual(len(indexes_1[1].parts), 1)
1293        self.assertTrue(isinstance(indexes_1[2], IncludeIndex))
1294        self.assertEqual(indexes_1[2].name, 'GenderIndex')
1295        self.assertEqual(len(indexes_1[2].parts), 2)
1296        self.assertEqual(indexes_1[2].includes_fields, ['gender'])
1297
1298        raw_indexes_2 = [
1299            {
1300                "IndexName": "MostRecentlyJoinedIndex",
1301                "KeySchema": [
1302                    {
1303                        "AttributeName": "username",
1304                        "KeyType": "HASH"
1305                    },
1306                    {
1307                        "AttributeName": "date_joined",
1308                        "KeyType": "RANGE"
1309                    }
1310                ],
1311                "Projection": {
1312                    "ProjectionType": "SOMETHING_CRAZY"
1313                }
1314            },
1315        ]
1316        self.assertRaises(
1317            exceptions.UnknownIndexFieldError,
1318            self.users._introspect_indexes,
1319            raw_indexes_2
1320        )
1321
1322    def test_initialization(self):
1323        users = Table('users', connection=self.default_connection)
1324        self.assertEqual(users.table_name, 'users')
1325        self.assertTrue(isinstance(users.connection, DynamoDBConnection))
1326        self.assertEqual(users.throughput['read'], 5)
1327        self.assertEqual(users.throughput['write'], 5)
1328        self.assertEqual(users.schema, None)
1329        self.assertEqual(users.indexes, None)
1330
1331        groups = Table('groups', connection=FakeDynamoDBConnection())
1332        self.assertEqual(groups.table_name, 'groups')
1333        self.assertTrue(hasattr(groups.connection, 'assert_called_once_with'))
1334
1335    def test_create_simple(self):
1336        conn = FakeDynamoDBConnection()
1337
1338        with mock.patch.object(conn, 'create_table', return_value={}) \
1339                as mock_create_table:
1340            retval = Table.create('users', schema=[
1341                HashKey('username'),
1342                RangeKey('date_joined', data_type=NUMBER)
1343            ], connection=conn)
1344            self.assertTrue(retval)
1345
1346        self.assertTrue(mock_create_table.called)
1347        mock_create_table.assert_called_once_with(attribute_definitions=[
1348            {
1349                'AttributeName': 'username',
1350                'AttributeType': 'S'
1351            },
1352            {
1353                'AttributeName': 'date_joined',
1354                'AttributeType': 'N'
1355            }
1356        ],
1357        table_name='users',
1358        key_schema=[
1359            {
1360                'KeyType': 'HASH',
1361                'AttributeName': 'username'
1362            },
1363            {
1364                'KeyType': 'RANGE',
1365                'AttributeName': 'date_joined'
1366            }
1367        ],
1368        provisioned_throughput={
1369            'WriteCapacityUnits': 5,
1370            'ReadCapacityUnits': 5
1371        })
1372
1373    def test_create_full(self):
1374        conn = FakeDynamoDBConnection()
1375
1376        with mock.patch.object(conn, 'create_table', return_value={}) \
1377                as mock_create_table:
1378            retval = Table.create('users', schema=[
1379                HashKey('username'),
1380                RangeKey('date_joined', data_type=NUMBER)
1381            ], throughput={
1382                'read':20,
1383                'write': 10,
1384            }, indexes=[
1385                KeysOnlyIndex('FriendCountIndex', parts=[
1386                    RangeKey('friend_count')
1387                ]),
1388            ], global_indexes=[
1389                GlobalKeysOnlyIndex('FullFriendCountIndex', parts=[
1390                    RangeKey('friend_count')
1391                ], throughput={
1392                    'read': 10,
1393                    'write': 8,
1394                }),
1395            ], connection=conn)
1396            self.assertTrue(retval)
1397
1398        self.assertTrue(mock_create_table.called)
1399        mock_create_table.assert_called_once_with(attribute_definitions=[
1400            {
1401                'AttributeName': 'username',
1402                'AttributeType': 'S'
1403            },
1404            {
1405                'AttributeName': 'date_joined',
1406                'AttributeType': 'N'
1407            },
1408            {
1409                'AttributeName': 'friend_count',
1410                'AttributeType': 'S'
1411            }
1412        ],
1413        key_schema=[
1414            {
1415                'KeyType': 'HASH',
1416                'AttributeName': 'username'
1417            },
1418            {
1419                'KeyType': 'RANGE',
1420                'AttributeName': 'date_joined'
1421            }
1422        ],
1423        table_name='users',
1424        provisioned_throughput={
1425            'WriteCapacityUnits': 10,
1426            'ReadCapacityUnits': 20
1427        },
1428        global_secondary_indexes=[
1429            {
1430                'KeySchema': [
1431                    {
1432                        'KeyType': 'RANGE',
1433                        'AttributeName': 'friend_count'
1434                    }
1435                ],
1436                'IndexName': 'FullFriendCountIndex',
1437                'Projection': {
1438                    'ProjectionType': 'KEYS_ONLY'
1439                },
1440                'ProvisionedThroughput': {
1441                    'WriteCapacityUnits': 8,
1442                    'ReadCapacityUnits': 10
1443                }
1444            }
1445        ],
1446        local_secondary_indexes=[
1447            {
1448                'KeySchema': [
1449                    {
1450                        'KeyType': 'RANGE',
1451                        'AttributeName': 'friend_count'
1452                    }
1453                ],
1454                'IndexName': 'FriendCountIndex',
1455                'Projection': {
1456                    'ProjectionType': 'KEYS_ONLY'
1457                }
1458            }
1459        ])
1460
1461    def test_describe(self):
1462        expected = {
1463            "Table": {
1464                "AttributeDefinitions": [
1465                    {
1466                        "AttributeName": "username",
1467                        "AttributeType": "S"
1468                    }
1469                ],
1470                "ItemCount": 5,
1471                "KeySchema": [
1472                    {
1473                        "AttributeName": "username",
1474                        "KeyType": "HASH"
1475                    }
1476                ],
1477                "LocalSecondaryIndexes": [
1478                    {
1479                        "IndexName": "UsernameIndex",
1480                        "KeySchema": [
1481                            {
1482                                "AttributeName": "username",
1483                                "KeyType": "HASH"
1484                            }
1485                        ],
1486                        "Projection": {
1487                            "ProjectionType": "KEYS_ONLY"
1488                        }
1489                    }
1490                ],
1491                "ProvisionedThroughput": {
1492                    "ReadCapacityUnits": 20,
1493                    "WriteCapacityUnits": 6
1494                },
1495                "TableName": "Thread",
1496                "TableStatus": "ACTIVE"
1497            }
1498        }
1499
1500        with mock.patch.object(
1501                self.users.connection,
1502                'describe_table',
1503                return_value=expected) as mock_describe:
1504            self.assertEqual(self.users.throughput['read'], 5)
1505            self.assertEqual(self.users.throughput['write'], 5)
1506            self.assertEqual(self.users.schema, None)
1507            self.assertEqual(self.users.indexes, None)
1508
1509            self.users.describe()
1510
1511            self.assertEqual(self.users.throughput['read'], 20)
1512            self.assertEqual(self.users.throughput['write'], 6)
1513            self.assertEqual(len(self.users.schema), 1)
1514            self.assertEqual(isinstance(self.users.schema[0], HashKey), 1)
1515            self.assertEqual(len(self.users.indexes), 1)
1516
1517        mock_describe.assert_called_once_with('users')
1518
1519    def test_update(self):
1520        with mock.patch.object(
1521                self.users.connection,
1522                'update_table',
1523                return_value={}) as mock_update:
1524            self.assertEqual(self.users.throughput['read'], 5)
1525            self.assertEqual(self.users.throughput['write'], 5)
1526            self.users.update(throughput={
1527                'read': 7,
1528                'write': 2,
1529            })
1530            self.assertEqual(self.users.throughput['read'], 7)
1531            self.assertEqual(self.users.throughput['write'], 2)
1532
1533        mock_update.assert_called_once_with(
1534            'users',
1535            global_secondary_index_updates=None,
1536            provisioned_throughput={
1537                'WriteCapacityUnits': 2,
1538                'ReadCapacityUnits': 7
1539            }
1540        )
1541
1542        with mock.patch.object(
1543                self.users.connection,
1544                'update_table',
1545                return_value={}) as mock_update:
1546            self.assertEqual(self.users.throughput['read'], 7)
1547            self.assertEqual(self.users.throughput['write'], 2)
1548            self.users.update(throughput={
1549                'read': 9,
1550                'write': 5,
1551            },
1552            global_indexes={
1553                'WhateverIndex': {
1554                    'read': 6,
1555                    'write': 1
1556                },
1557                'AnotherIndex': {
1558                    'read': 1,
1559                    'write': 2
1560                }
1561            })
1562            self.assertEqual(self.users.throughput['read'], 9)
1563            self.assertEqual(self.users.throughput['write'], 5)
1564
1565        args, kwargs = mock_update.call_args
1566        self.assertEqual(args, ('users',))
1567        self.assertEqual(kwargs['provisioned_throughput'], {
1568            'WriteCapacityUnits': 5,
1569            'ReadCapacityUnits': 9,
1570            })
1571        update = kwargs['global_secondary_index_updates'][:]
1572        update.sort(key=lambda x: x['Update']['IndexName'])
1573        self.assertDictEqual(
1574            update[0],
1575            {
1576                'Update': {
1577                    'IndexName': 'AnotherIndex',
1578                    'ProvisionedThroughput': {
1579                        'WriteCapacityUnits': 2,
1580                        'ReadCapacityUnits': 1
1581                    }
1582                }
1583            })
1584        self.assertDictEqual(
1585            update[1],
1586            {
1587                'Update': {
1588                    'IndexName': 'WhateverIndex',
1589                    'ProvisionedThroughput': {
1590                        'WriteCapacityUnits': 1,
1591                        'ReadCapacityUnits': 6
1592                    }
1593                }
1594            })
1595
1596    def test_create_global_secondary_index(self):
1597        with mock.patch.object(
1598                self.users.connection,
1599                'update_table',
1600                return_value={}) as mock_update:
1601            self.users.create_global_secondary_index(
1602                global_index=GlobalAllIndex(
1603                    'JustCreatedIndex',
1604                    parts=[
1605                        HashKey('requiredHashKey')
1606                    ],
1607                    throughput={
1608                        'read': 2,
1609                        'write': 2
1610                    }
1611                )
1612            )
1613
1614        mock_update.assert_called_once_with(
1615            'users',
1616            global_secondary_index_updates=[
1617                {
1618                    'Create': {
1619                        'IndexName': 'JustCreatedIndex',
1620                        'KeySchema': [
1621                            {
1622                                'KeyType': 'HASH',
1623                                'AttributeName': 'requiredHashKey'
1624                            }
1625                        ],
1626                        'Projection': {
1627                            'ProjectionType': 'ALL'
1628                        },
1629                        'ProvisionedThroughput': {
1630                            'WriteCapacityUnits': 2,
1631                            'ReadCapacityUnits': 2
1632                        }
1633                    }
1634                }
1635            ],
1636            attribute_definitions=[
1637                {
1638                    'AttributeName': 'requiredHashKey',
1639                    'AttributeType': 'S'
1640                }
1641            ]
1642        )
1643
1644    def test_delete_global_secondary_index(self):
1645        with mock.patch.object(
1646                self.users.connection,
1647                'update_table',
1648                return_value={}) as mock_update:
1649            self.users.delete_global_secondary_index('RandomGSIIndex')
1650
1651        mock_update.assert_called_once_with(
1652            'users',
1653            global_secondary_index_updates=[
1654                {
1655                    'Delete': {
1656                        'IndexName': 'RandomGSIIndex',
1657                    }
1658                }
1659            ]
1660        )
1661
1662    def test_update_global_secondary_index(self):
1663        # Updating a single global secondary index
1664        with mock.patch.object(
1665                self.users.connection,
1666                'update_table',
1667                return_value={}) as mock_update:
1668            self.users.update_global_secondary_index(global_indexes={
1669                'A_IndexToBeUpdated': {
1670                    'read': 5,
1671                    'write': 5
1672                }
1673            })
1674
1675        mock_update.assert_called_once_with(
1676            'users',
1677            global_secondary_index_updates=[
1678                {
1679                    'Update': {
1680                        'IndexName': 'A_IndexToBeUpdated',
1681                        "ProvisionedThroughput": {
1682                            "ReadCapacityUnits": 5,
1683                            "WriteCapacityUnits": 5
1684                        },
1685                    }
1686                }
1687            ]
1688        )
1689
1690        # Updating multiple global secondary indexes
1691        with mock.patch.object(
1692                self.users.connection,
1693                'update_table',
1694                return_value={}) as mock_update:
1695            self.users.update_global_secondary_index(global_indexes={
1696                'A_IndexToBeUpdated': {
1697                    'read': 5,
1698                    'write': 5
1699                },
1700                'B_IndexToBeUpdated': {
1701                    'read': 9,
1702                    'write': 9
1703                }
1704            })
1705
1706        args, kwargs = mock_update.call_args
1707        self.assertEqual(args, ('users',))
1708        update = kwargs['global_secondary_index_updates'][:]
1709        update.sort(key=lambda x: x['Update']['IndexName'])
1710        self.assertDictEqual(
1711            update[0],
1712            {
1713                'Update': {
1714                    'IndexName': 'A_IndexToBeUpdated',
1715                    'ProvisionedThroughput': {
1716                        'WriteCapacityUnits': 5,
1717                        'ReadCapacityUnits': 5
1718                    }
1719                }
1720            })
1721        self.assertDictEqual(
1722            update[1],
1723            {
1724                'Update': {
1725                    'IndexName': 'B_IndexToBeUpdated',
1726                    'ProvisionedThroughput': {
1727                        'WriteCapacityUnits': 9,
1728                        'ReadCapacityUnits': 9
1729                    }
1730                }
1731            })
1732
1733    def test_delete(self):
1734        with mock.patch.object(
1735                self.users.connection,
1736                'delete_table',
1737                return_value={}) as mock_delete:
1738            self.assertTrue(self.users.delete())
1739
1740        mock_delete.assert_called_once_with('users')
1741
1742    def test_get_item(self):
1743        expected = {
1744            'Item': {
1745                'username': {'S': 'johndoe'},
1746                'first_name': {'S': 'John'},
1747                'last_name': {'S': 'Doe'},
1748                'date_joined': {'N': '1366056668'},
1749                'friend_count': {'N': '3'},
1750                'friends': {'SS': ['alice', 'bob', 'jane']},
1751            }
1752        }
1753
1754        with mock.patch.object(
1755                self.users.connection,
1756                'get_item',
1757                return_value=expected) as mock_get_item:
1758            item = self.users.get_item(username='johndoe')
1759            self.assertEqual(item['username'], 'johndoe')
1760            self.assertEqual(item['first_name'], 'John')
1761
1762        mock_get_item.assert_called_once_with('users', {
1763            'username': {'S': 'johndoe'}
1764        }, consistent_read=False, attributes_to_get=None)
1765
1766        with mock.patch.object(
1767                self.users.connection,
1768                'get_item',
1769                return_value=expected) as mock_get_item:
1770            item = self.users.get_item(username='johndoe', attributes=[
1771                'username',
1772                'first_name',
1773            ])
1774
1775        mock_get_item.assert_called_once_with('users', {
1776            'username': {'S': 'johndoe'}
1777        }, consistent_read=False, attributes_to_get=['username', 'first_name'])
1778
1779    def test_has_item(self):
1780        expected = {
1781            'Item': {
1782                'username': {'S': 'johndoe'},
1783                'first_name': {'S': 'John'},
1784                'last_name': {'S': 'Doe'},
1785                'date_joined': {'N': '1366056668'},
1786                'friend_count': {'N': '3'},
1787                'friends': {'SS': ['alice', 'bob', 'jane']},
1788            }
1789        }
1790
1791        with mock.patch.object(
1792                self.users.connection,
1793                'get_item',
1794                return_value=expected) as mock_get_item:
1795            found = self.users.has_item(username='johndoe')
1796            self.assertTrue(found)
1797
1798        with mock.patch.object(
1799                self.users.connection,
1800                'get_item') as mock_get_item:
1801            mock_get_item.side_effect = JSONResponseError("Nope.", None, None)
1802            found = self.users.has_item(username='mrsmith')
1803            self.assertFalse(found)
1804
1805    def test_lookup_hash(self):
1806        """Tests the "lookup" function with just a hash key"""
1807        expected = {
1808            'Item': {
1809                'username': {'S': 'johndoe'},
1810                'first_name': {'S': 'John'},
1811                'last_name': {'S': 'Doe'},
1812                'date_joined': {'N': '1366056668'},
1813                'friend_count': {'N': '3'},
1814                'friends': {'SS': ['alice', 'bob', 'jane']},
1815            }
1816        }
1817
1818        # Set the Schema
1819        self.users.schema = [
1820            HashKey('username'),
1821            RangeKey('date_joined', data_type=NUMBER),
1822        ]
1823
1824        with mock.patch.object(
1825                self.users,
1826                'get_item',
1827                return_value=expected) as mock_get_item:
1828            self.users.lookup('johndoe')
1829
1830        mock_get_item.assert_called_once_with(
1831            username= 'johndoe')
1832
1833    def test_lookup_hash_and_range(self):
1834        """Test the "lookup" function with a hash and range key"""
1835        expected = {
1836            'Item': {
1837                'username': {'S': 'johndoe'},
1838                'first_name': {'S': 'John'},
1839                'last_name': {'S': 'Doe'},
1840                'date_joined': {'N': '1366056668'},
1841                'friend_count': {'N': '3'},
1842                'friends': {'SS': ['alice', 'bob', 'jane']},
1843            }
1844        }
1845
1846        # Set the Schema
1847        self.users.schema = [
1848            HashKey('username'),
1849            RangeKey('date_joined', data_type=NUMBER),
1850        ]
1851
1852        with mock.patch.object(
1853                self.users,
1854                'get_item',
1855                return_value=expected) as mock_get_item:
1856            self.users.lookup('johndoe', 1366056668)
1857
1858        mock_get_item.assert_called_once_with(
1859            username= 'johndoe',
1860            date_joined= 1366056668)
1861
1862    def test_put_item(self):
1863        with mock.patch.object(
1864                self.users.connection,
1865                'put_item',
1866                return_value={}) as mock_put_item:
1867            self.users.put_item(data={
1868                'username': 'johndoe',
1869                'last_name': 'Doe',
1870                'date_joined': 12345,
1871            })
1872
1873        mock_put_item.assert_called_once_with('users', {
1874            'username': {'S': 'johndoe'},
1875            'last_name': {'S': 'Doe'},
1876            'date_joined': {'N': '12345'}
1877        }, expected={
1878            'username': {
1879                'Exists': False,
1880            },
1881            'last_name': {
1882                'Exists': False,
1883            },
1884            'date_joined': {
1885                'Exists': False,
1886            }
1887        })
1888
1889    def test_private_put_item(self):
1890        with mock.patch.object(
1891                self.users.connection,
1892                'put_item',
1893                return_value={}) as mock_put_item:
1894            self.users._put_item({'some': 'data'})
1895
1896        mock_put_item.assert_called_once_with('users', {'some': 'data'})
1897
1898    def test_private_update_item(self):
1899        with mock.patch.object(
1900                self.users.connection,
1901                'update_item',
1902                return_value={}) as mock_update_item:
1903            self.users._update_item({
1904                'username': 'johndoe'
1905            }, {
1906                'some': 'data',
1907            })
1908
1909        mock_update_item.assert_called_once_with('users', {
1910            'username': {'S': 'johndoe'},
1911        }, {
1912            'some': 'data',
1913        })
1914
1915    def test_delete_item(self):
1916        with mock.patch.object(
1917                self.users.connection,
1918                'delete_item',
1919                return_value={}) as mock_delete_item:
1920            self.assertTrue(self.users.delete_item(username='johndoe', date_joined=23456))
1921
1922        mock_delete_item.assert_called_once_with('users', {
1923            'username': {
1924                'S': 'johndoe'
1925            },
1926            'date_joined': {
1927                'N': '23456'
1928            }
1929        }, expected=None, conditional_operator=None)
1930
1931    def test_delete_item_conditionally(self):
1932        with mock.patch.object(
1933                self.users.connection,
1934                'delete_item',
1935                return_value={}) as mock_delete_item:
1936            self.assertTrue(self.users.delete_item(expected={'balance__eq': 0},
1937                                                   username='johndoe', date_joined=23456))
1938
1939        mock_delete_item.assert_called_once_with('users', {
1940            'username': {
1941                'S': 'johndoe'
1942            },
1943            'date_joined': {
1944                'N': '23456'
1945            }
1946        },
1947        expected={
1948            'balance': {
1949                'ComparisonOperator': 'EQ', 'AttributeValueList': [{'N': '0'}]
1950                },
1951            },
1952        conditional_operator=None)
1953
1954        def side_effect(*args, **kwargs):
1955            raise exceptions.ConditionalCheckFailedException(400, '', {})
1956
1957        with mock.patch.object(
1958                self.users.connection,
1959                'delete_item',
1960                side_effect=side_effect) as mock_delete_item:
1961            self.assertFalse(self.users.delete_item(expected={'balance__eq': 0},
1962                                                    username='johndoe', date_joined=23456))
1963
1964    def test_get_key_fields_no_schema_populated(self):
1965        expected = {
1966            "Table": {
1967                "AttributeDefinitions": [
1968                    {
1969                        "AttributeName": "username",
1970                        "AttributeType": "S"
1971                    },
1972                    {
1973                        "AttributeName": "date_joined",
1974                        "AttributeType": "N"
1975                    }
1976                ],
1977                "ItemCount": 5,
1978                "KeySchema": [
1979                    {
1980                        "AttributeName": "username",
1981                        "KeyType": "HASH"
1982                    },
1983                    {
1984                        "AttributeName": "date_joined",
1985                        "KeyType": "RANGE"
1986                    }
1987                ],
1988                "LocalSecondaryIndexes": [
1989                    {
1990                        "IndexName": "UsernameIndex",
1991                        "KeySchema": [
1992                            {
1993                                "AttributeName": "username",
1994                                "KeyType": "HASH"
1995                            }
1996                        ],
1997                        "Projection": {
1998                            "ProjectionType": "KEYS_ONLY"
1999                        }
2000                    }
2001                ],
2002                "ProvisionedThroughput": {
2003                    "ReadCapacityUnits": 20,
2004                    "WriteCapacityUnits": 6
2005                },
2006                "TableName": "Thread",
2007                "TableStatus": "ACTIVE"
2008            }
2009        }
2010
2011        with mock.patch.object(
2012                self.users.connection,
2013                'describe_table',
2014                return_value=expected) as mock_describe:
2015            self.assertEqual(self.users.schema, None)
2016
2017            key_fields = self.users.get_key_fields()
2018            self.assertEqual(key_fields, ['username', 'date_joined'])
2019
2020            self.assertEqual(len(self.users.schema), 2)
2021
2022        mock_describe.assert_called_once_with('users')
2023
2024    def test_batch_write_no_writes(self):
2025        with mock.patch.object(
2026                self.users.connection,
2027                'batch_write_item',
2028                return_value={}) as mock_batch:
2029            with self.users.batch_write() as batch:
2030                pass
2031
2032        self.assertFalse(mock_batch.called)
2033
2034    def test_batch_write(self):
2035        with mock.patch.object(
2036                self.users.connection,
2037                'batch_write_item',
2038                return_value={}) as mock_batch:
2039            with self.users.batch_write() as batch:
2040                batch.put_item(data={
2041                    'username': 'jane',
2042                    'date_joined': 12342547
2043                })
2044                batch.delete_item(username='johndoe')
2045                batch.put_item(data={
2046                    'username': 'alice',
2047                    'date_joined': 12342888
2048                })
2049
2050        mock_batch.assert_called_once_with({
2051            'users': [
2052                {
2053                    'PutRequest': {
2054                        'Item': {
2055                            'username': {'S': 'jane'},
2056                            'date_joined': {'N': '12342547'}
2057                        }
2058                    }
2059                },
2060                {
2061                    'PutRequest': {
2062                        'Item': {
2063                            'username': {'S': 'alice'},
2064                            'date_joined': {'N': '12342888'}
2065                        }
2066                    }
2067                },
2068                {
2069                    'DeleteRequest': {
2070                        'Key': {
2071                            'username': {'S': 'johndoe'},
2072                        }
2073                    }
2074                },
2075            ]
2076        })
2077
2078    def test_batch_write_dont_swallow_exceptions(self):
2079        with mock.patch.object(
2080                self.users.connection,
2081                'batch_write_item',
2082                return_value={}) as mock_batch:
2083            try:
2084                with self.users.batch_write() as batch:
2085                    raise Exception('OH NOES')
2086            except Exception as e:
2087                self.assertEqual(str(e), 'OH NOES')
2088
2089        self.assertFalse(mock_batch.called)
2090
2091    def test_batch_write_flushing(self):
2092        with mock.patch.object(
2093                self.users.connection,
2094                'batch_write_item',
2095                return_value={}) as mock_batch:
2096            with self.users.batch_write() as batch:
2097                batch.put_item(data={
2098                    'username': 'jane',
2099                    'date_joined': 12342547
2100                })
2101                # This would only be enough for one batch.
2102                batch.delete_item(username='johndoe1')
2103                batch.delete_item(username='johndoe2')
2104                batch.delete_item(username='johndoe3')
2105                batch.delete_item(username='johndoe4')
2106                batch.delete_item(username='johndoe5')
2107                batch.delete_item(username='johndoe6')
2108                batch.delete_item(username='johndoe7')
2109                batch.delete_item(username='johndoe8')
2110                batch.delete_item(username='johndoe9')
2111                batch.delete_item(username='johndoe10')
2112                batch.delete_item(username='johndoe11')
2113                batch.delete_item(username='johndoe12')
2114                batch.delete_item(username='johndoe13')
2115                batch.delete_item(username='johndoe14')
2116                batch.delete_item(username='johndoe15')
2117                batch.delete_item(username='johndoe16')
2118                batch.delete_item(username='johndoe17')
2119                batch.delete_item(username='johndoe18')
2120                batch.delete_item(username='johndoe19')
2121                batch.delete_item(username='johndoe20')
2122                batch.delete_item(username='johndoe21')
2123                batch.delete_item(username='johndoe22')
2124                batch.delete_item(username='johndoe23')
2125
2126                # We're only at 24 items. No flushing yet.
2127                self.assertEqual(mock_batch.call_count, 0)
2128
2129                # This pushes it over the edge. A flush happens then we start
2130                # queuing objects again.
2131                batch.delete_item(username='johndoe24')
2132                self.assertEqual(mock_batch.call_count, 1)
2133                # Since we add another, there's enough for a second call to
2134                # flush.
2135                batch.delete_item(username='johndoe25')
2136
2137        self.assertEqual(mock_batch.call_count, 2)
2138
2139    def test_batch_write_unprocessed_items(self):
2140        unprocessed = {
2141            'UnprocessedItems': {
2142                'users': [
2143                    {
2144                        'PutRequest': {
2145                            'username': {
2146                                'S': 'jane',
2147                            },
2148                            'date_joined': {
2149                                'N': 12342547
2150                            }
2151                        },
2152                    },
2153                ],
2154            },
2155        }
2156
2157        # Test enqueuing the unprocessed bits.
2158        with mock.patch.object(
2159                self.users.connection,
2160                'batch_write_item',
2161                return_value=unprocessed) as mock_batch:
2162            with self.users.batch_write() as batch:
2163                self.assertEqual(len(batch._unprocessed), 0)
2164
2165                # Trash the ``resend_unprocessed`` method so that we don't
2166                # infinite loop forever here.
2167                batch.resend_unprocessed = lambda: True
2168
2169                batch.put_item(data={
2170                    'username': 'jane',
2171                    'date_joined': 12342547
2172                })
2173                batch.delete_item(username='johndoe')
2174                batch.put_item(data={
2175                    'username': 'alice',
2176                    'date_joined': 12342888
2177                })
2178
2179            self.assertEqual(len(batch._unprocessed), 1)
2180
2181        # Now test resending those unprocessed items.
2182        with mock.patch.object(
2183                self.users.connection,
2184                'batch_write_item',
2185                return_value={}) as mock_batch:
2186            with self.users.batch_write() as batch:
2187                self.assertEqual(len(batch._unprocessed), 0)
2188
2189                # Toss in faked unprocessed items, as though a previous batch
2190                # had failed.
2191                batch._unprocessed = [
2192                    {
2193                        'PutRequest': {
2194                            'username': {
2195                                'S': 'jane',
2196                            },
2197                            'date_joined': {
2198                                'N': 12342547
2199                            }
2200                        },
2201                    },
2202                ]
2203
2204                batch.put_item(data={
2205                    'username': 'jane',
2206                    'date_joined': 12342547
2207                })
2208                batch.delete_item(username='johndoe')
2209                batch.put_item(data={
2210                    'username': 'alice',
2211                    'date_joined': 12342888
2212                })
2213
2214                # Flush, to make sure everything has been processed.
2215                # Unprocessed items should still be hanging around.
2216                batch.flush()
2217                self.assertEqual(len(batch._unprocessed), 1)
2218
2219            # Post-exit, this should be emptied.
2220            self.assertEqual(len(batch._unprocessed), 0)
2221
2222    def test__build_filters(self):
2223        filters = self.users._build_filters({
2224            'username__eq': 'johndoe',
2225            'date_joined__gte': 1234567,
2226            'age__in': [30, 31, 32, 33],
2227            'last_name__between': ['danzig', 'only'],
2228            'first_name__null': False,
2229            'gender__null': True,
2230        }, using=FILTER_OPERATORS)
2231        self.assertEqual(filters, {
2232            'username': {
2233                'AttributeValueList': [
2234                    {
2235                        'S': 'johndoe',
2236                    },
2237                ],
2238                'ComparisonOperator': 'EQ',
2239            },
2240            'date_joined': {
2241                'AttributeValueList': [
2242                    {
2243                        'N': '1234567',
2244                    },
2245                ],
2246                'ComparisonOperator': 'GE',
2247            },
2248            'age': {
2249                'AttributeValueList': [
2250                    {'N': '30'},
2251                    {'N': '31'},
2252                    {'N': '32'},
2253                    {'N': '33'},
2254                ],
2255                'ComparisonOperator': 'IN',
2256            },
2257            'last_name': {
2258                'AttributeValueList': [{'S': 'danzig'}, {'S': 'only'}],
2259                'ComparisonOperator': 'BETWEEN',
2260            },
2261            'first_name': {
2262                'ComparisonOperator': 'NOT_NULL'
2263            },
2264            'gender': {
2265                'ComparisonOperator': 'NULL'
2266            },
2267        })
2268
2269        self.assertRaises(exceptions.UnknownFilterTypeError,
2270            self.users._build_filters,
2271            {
2272                'darling__die': True,
2273            }
2274        )
2275
2276        q_filters = self.users._build_filters({
2277            'username__eq': 'johndoe',
2278            'date_joined__gte': 1234567,
2279            'last_name__between': ['danzig', 'only'],
2280            'gender__beginswith': 'm',
2281        }, using=QUERY_OPERATORS)
2282        self.assertEqual(q_filters, {
2283            'username': {
2284                'AttributeValueList': [
2285                    {
2286                        'S': 'johndoe',
2287                    },
2288                ],
2289                'ComparisonOperator': 'EQ',
2290            },
2291            'date_joined': {
2292                'AttributeValueList': [
2293                    {
2294                        'N': '1234567',
2295                    },
2296                ],
2297                'ComparisonOperator': 'GE',
2298            },
2299            'last_name': {
2300                'AttributeValueList': [{'S': 'danzig'}, {'S': 'only'}],
2301                'ComparisonOperator': 'BETWEEN',
2302            },
2303            'gender': {
2304                'AttributeValueList': [{'S': 'm'}],
2305                'ComparisonOperator': 'BEGINS_WITH',
2306            },
2307        })
2308
2309        self.assertRaises(exceptions.UnknownFilterTypeError,
2310            self.users._build_filters,
2311            {
2312                'darling__die': True,
2313            },
2314            using=QUERY_OPERATORS
2315        )
2316        self.assertRaises(exceptions.UnknownFilterTypeError,
2317            self.users._build_filters,
2318            {
2319                'first_name__null': True,
2320            },
2321            using=QUERY_OPERATORS
2322        )
2323
2324    def test_private_query(self):
2325        expected = {
2326            "ConsumedCapacity": {
2327                "CapacityUnits": 0.5,
2328                "TableName": "users"
2329            },
2330            "Count": 4,
2331            "Items": [
2332                {
2333                    'username': {'S': 'johndoe'},
2334                    'first_name': {'S': 'John'},
2335                    'last_name': {'S': 'Doe'},
2336                    'date_joined': {'N': '1366056668'},
2337                    'friend_count': {'N': '3'},
2338                    'friends': {'SS': ['alice', 'bob', 'jane']},
2339                },
2340                {
2341                    'username': {'S': 'jane'},
2342                    'first_name': {'S': 'Jane'},
2343                    'last_name': {'S': 'Doe'},
2344                    'date_joined': {'N': '1366057777'},
2345                    'friend_count': {'N': '2'},
2346                    'friends': {'SS': ['alice', 'johndoe']},
2347                },
2348                {
2349                    'username': {'S': 'alice'},
2350                    'first_name': {'S': 'Alice'},
2351                    'last_name': {'S': 'Expert'},
2352                    'date_joined': {'N': '1366056680'},
2353                    'friend_count': {'N': '1'},
2354                    'friends': {'SS': ['jane']},
2355                },
2356                {
2357                    'username': {'S': 'bob'},
2358                    'first_name': {'S': 'Bob'},
2359                    'last_name': {'S': 'Smith'},
2360                    'date_joined': {'N': '1366056888'},
2361                    'friend_count': {'N': '1'},
2362                    'friends': {'SS': ['johndoe']},
2363                },
2364            ],
2365            "ScannedCount": 4
2366        }
2367
2368        with mock.patch.object(
2369                self.users.connection,
2370                'query',
2371                return_value=expected) as mock_query:
2372            results = self.users._query(
2373                limit=4,
2374                reverse=True,
2375                username__between=['aaa', 'mmm']
2376            )
2377            usernames = [res['username'] for res in results['results']]
2378            self.assertEqual(usernames, ['johndoe', 'jane', 'alice', 'bob'])
2379            self.assertEqual(len(results['results']), 4)
2380            self.assertEqual(results['last_key'], None)
2381
2382        mock_query.assert_called_once_with('users',
2383            consistent_read=False,
2384            scan_index_forward=False,
2385            index_name=None,
2386            attributes_to_get=None,
2387            limit=4,
2388            key_conditions={
2389                'username': {
2390                    'AttributeValueList': [{'S': 'aaa'}, {'S': 'mmm'}],
2391                    'ComparisonOperator': 'BETWEEN',
2392                }
2393            },
2394            select=None,
2395            query_filter=None,
2396            conditional_operator=None
2397        )
2398
2399        # Now alter the expected.
2400        expected['LastEvaluatedKey'] = {
2401            'username': {
2402                'S': 'johndoe',
2403            },
2404        }
2405
2406        with mock.patch.object(
2407                self.users.connection,
2408                'query',
2409                return_value=expected) as mock_query_2:
2410            results = self.users._query(
2411                limit=4,
2412                reverse=True,
2413                username__between=['aaa', 'mmm'],
2414                exclusive_start_key={
2415                    'username': 'adam',
2416                },
2417                consistent=True,
2418                query_filter=None,
2419                conditional_operator='AND'
2420            )
2421            usernames = [res['username'] for res in results['results']]
2422            self.assertEqual(usernames, ['johndoe', 'jane', 'alice', 'bob'])
2423            self.assertEqual(len(results['results']), 4)
2424            self.assertEqual(results['last_key'], {'username': 'johndoe'})
2425
2426        mock_query_2.assert_called_once_with('users',
2427            key_conditions={
2428                'username': {
2429                    'AttributeValueList': [{'S': 'aaa'}, {'S': 'mmm'}],
2430                    'ComparisonOperator': 'BETWEEN',
2431                }
2432            },
2433            index_name=None,
2434            attributes_to_get=None,
2435            scan_index_forward=False,
2436            limit=4,
2437            exclusive_start_key={
2438                'username': {
2439                    'S': 'adam',
2440                },
2441            },
2442            consistent_read=True,
2443            select=None,
2444            query_filter=None,
2445            conditional_operator='AND'
2446        )
2447
2448    def test_private_scan(self):
2449        expected = {
2450            "ConsumedCapacity": {
2451                "CapacityUnits": 0.5,
2452                "TableName": "users"
2453            },
2454            "Count": 4,
2455            "Items": [
2456                {
2457                    'username': {'S': 'alice'},
2458                    'first_name': {'S': 'Alice'},
2459                    'last_name': {'S': 'Expert'},
2460                    'date_joined': {'N': '1366056680'},
2461                    'friend_count': {'N': '1'},
2462                    'friends': {'SS': ['jane']},
2463                },
2464                {
2465                    'username': {'S': 'bob'},
2466                    'first_name': {'S': 'Bob'},
2467                    'last_name': {'S': 'Smith'},
2468                    'date_joined': {'N': '1366056888'},
2469                    'friend_count': {'N': '1'},
2470                    'friends': {'SS': ['johndoe']},
2471                },
2472                {
2473                    'username': {'S': 'jane'},
2474                    'first_name': {'S': 'Jane'},
2475                    'last_name': {'S': 'Doe'},
2476                    'date_joined': {'N': '1366057777'},
2477                    'friend_count': {'N': '2'},
2478                    'friends': {'SS': ['alice', 'johndoe']},
2479                },
2480            ],
2481            "ScannedCount": 4
2482        }
2483
2484        with mock.patch.object(
2485                self.users.connection,
2486                'scan',
2487                return_value=expected) as mock_scan:
2488            results = self.users._scan(
2489                limit=2,
2490                friend_count__lte=2
2491            )
2492            usernames = [res['username'] for res in results['results']]
2493            self.assertEqual(usernames, ['alice', 'bob', 'jane'])
2494            self.assertEqual(len(results['results']), 3)
2495            self.assertEqual(results['last_key'], None)
2496
2497        mock_scan.assert_called_once_with('users',
2498            scan_filter={
2499                'friend_count': {
2500                    'AttributeValueList': [{'N': '2'}],
2501                    'ComparisonOperator': 'LE',
2502                }
2503            },
2504            limit=2,
2505            segment=None,
2506            attributes_to_get=None,
2507            total_segments=None,
2508            conditional_operator=None
2509        )
2510
2511        # Now alter the expected.
2512        expected['LastEvaluatedKey'] = {
2513            'username': {
2514                'S': 'jane',
2515            },
2516        }
2517
2518        with mock.patch.object(
2519                self.users.connection,
2520                'scan',
2521                return_value=expected) as mock_scan_2:
2522            results = self.users._scan(
2523                limit=3,
2524                friend_count__lte=2,
2525                exclusive_start_key={
2526                    'username': 'adam',
2527                },
2528                segment=None,
2529                total_segments=None
2530            )
2531            usernames = [res['username'] for res in results['results']]
2532            self.assertEqual(usernames, ['alice', 'bob', 'jane'])
2533            self.assertEqual(len(results['results']), 3)
2534            self.assertEqual(results['last_key'], {'username': 'jane'})
2535
2536        mock_scan_2.assert_called_once_with('users',
2537            scan_filter={
2538                'friend_count': {
2539                    'AttributeValueList': [{'N': '2'}],
2540                    'ComparisonOperator': 'LE',
2541                }
2542            },
2543            limit=3,
2544            exclusive_start_key={
2545                'username': {
2546                    'S': 'adam',
2547                },
2548            },
2549            segment=None,
2550            attributes_to_get=None,
2551            total_segments=None,
2552            conditional_operator=None
2553        )
2554
2555    def test_query(self):
2556        items_1 = {
2557            'results': [
2558                Item(self.users, data={
2559                    'username': 'johndoe',
2560                    'first_name': 'John',
2561                    'last_name': 'Doe',
2562                }),
2563                Item(self.users, data={
2564                    'username': 'jane',
2565                    'first_name': 'Jane',
2566                    'last_name': 'Doe',
2567                }),
2568            ],
2569            'last_key': 'jane',
2570        }
2571
2572        results = self.users.query_2(last_name__eq='Doe')
2573        self.assertTrue(isinstance(results, ResultSet))
2574        self.assertEqual(len(results._results), 0)
2575        self.assertEqual(results.the_callable, self.users._query)
2576
2577        with mock.patch.object(
2578                results,
2579                'the_callable',
2580                return_value=items_1) as mock_query:
2581            res_1 = next(results)
2582            # Now it should be populated.
2583            self.assertEqual(len(results._results), 2)
2584            self.assertEqual(res_1['username'], 'johndoe')
2585            res_2 = next(results)
2586            self.assertEqual(res_2['username'], 'jane')
2587
2588        self.assertEqual(mock_query.call_count, 1)
2589
2590        items_2 = {
2591            'results': [
2592                Item(self.users, data={
2593                    'username': 'foodoe',
2594                    'first_name': 'Foo',
2595                    'last_name': 'Doe',
2596                }),
2597            ],
2598        }
2599
2600        with mock.patch.object(
2601                results,
2602                'the_callable',
2603                return_value=items_2) as mock_query_2:
2604            res_3 = next(results)
2605            # New results should have been found.
2606            self.assertEqual(len(results._results), 1)
2607            self.assertEqual(res_3['username'], 'foodoe')
2608
2609            self.assertRaises(StopIteration, results.next)
2610
2611        self.assertEqual(mock_query_2.call_count, 1)
2612
2613    def test_query_with_specific_attributes(self):
2614        items_1 = {
2615            'results': [
2616                Item(self.users, data={
2617                    'username': 'johndoe',
2618                }),
2619                Item(self.users, data={
2620                    'username': 'jane',
2621                }),
2622            ],
2623            'last_key': 'jane',
2624        }
2625
2626        results = self.users.query_2(last_name__eq='Doe',
2627                                   attributes=['username'])
2628        self.assertTrue(isinstance(results, ResultSet))
2629        self.assertEqual(len(results._results), 0)
2630        self.assertEqual(results.the_callable, self.users._query)
2631
2632        with mock.patch.object(
2633                results,
2634                'the_callable',
2635                return_value=items_1) as mock_query:
2636            res_1 = next(results)
2637            # Now it should be populated.
2638            self.assertEqual(len(results._results), 2)
2639            self.assertEqual(res_1['username'], 'johndoe')
2640            self.assertEqual(list(res_1.keys()), ['username'])
2641            res_2 = next(results)
2642            self.assertEqual(res_2['username'], 'jane')
2643
2644        self.assertEqual(mock_query.call_count, 1)
2645
2646    def test_scan(self):
2647        items_1 = {
2648            'results': [
2649                Item(self.users, data={
2650                    'username': 'johndoe',
2651                    'first_name': 'John',
2652                    'last_name': 'Doe',
2653                }),
2654                Item(self.users, data={
2655                    'username': 'jane',
2656                    'first_name': 'Jane',
2657                    'last_name': 'Doe',
2658                }),
2659            ],
2660            'last_key': 'jane',
2661        }
2662
2663        results = self.users.scan(last_name__eq='Doe')
2664        self.assertTrue(isinstance(results, ResultSet))
2665        self.assertEqual(len(results._results), 0)
2666        self.assertEqual(results.the_callable, self.users._scan)
2667
2668        with mock.patch.object(
2669                results,
2670                'the_callable',
2671                return_value=items_1) as mock_scan:
2672            res_1 = next(results)
2673            # Now it should be populated.
2674            self.assertEqual(len(results._results), 2)
2675            self.assertEqual(res_1['username'], 'johndoe')
2676            res_2 = next(results)
2677            self.assertEqual(res_2['username'], 'jane')
2678
2679        self.assertEqual(mock_scan.call_count, 1)
2680
2681        items_2 = {
2682            'results': [
2683                Item(self.users, data={
2684                    'username': 'zoeydoe',
2685                    'first_name': 'Zoey',
2686                    'last_name': 'Doe',
2687                }),
2688            ],
2689        }
2690
2691        with mock.patch.object(
2692                results,
2693                'the_callable',
2694                return_value=items_2) as mock_scan_2:
2695            res_3 = next(results)
2696            # New results should have been found.
2697            self.assertEqual(len(results._results), 1)
2698            self.assertEqual(res_3['username'], 'zoeydoe')
2699
2700            self.assertRaises(StopIteration, results.next)
2701
2702        self.assertEqual(mock_scan_2.call_count, 1)
2703
2704    def test_scan_with_specific_attributes(self):
2705        items_1 = {
2706            'results': [
2707                Item(self.users, data={
2708                    'username': 'johndoe',
2709                }),
2710                Item(self.users, data={
2711                    'username': 'jane',
2712                }),
2713            ],
2714            'last_key': 'jane',
2715        }
2716
2717        results = self.users.scan(attributes=['username'])
2718        self.assertTrue(isinstance(results, ResultSet))
2719        self.assertEqual(len(results._results), 0)
2720        self.assertEqual(results.the_callable, self.users._scan)
2721
2722        with mock.patch.object(
2723                results,
2724                'the_callable',
2725                return_value=items_1) as mock_query:
2726            res_1 = next(results)
2727            # Now it should be populated.
2728            self.assertEqual(len(results._results), 2)
2729            self.assertEqual(res_1['username'], 'johndoe')
2730            self.assertEqual(list(res_1.keys()), ['username'])
2731            res_2 = next(results)
2732            self.assertEqual(res_2['username'], 'jane')
2733
2734        self.assertEqual(mock_query.call_count, 1)
2735
2736    def test_count(self):
2737        expected = {
2738            "Table": {
2739                "AttributeDefinitions": [
2740                    {
2741                        "AttributeName": "username",
2742                        "AttributeType": "S"
2743                    }
2744                ],
2745                "ItemCount": 5,
2746                "KeySchema": [
2747                    {
2748                        "AttributeName": "username",
2749                        "KeyType": "HASH"
2750                    }
2751                ],
2752                "LocalSecondaryIndexes": [
2753                    {
2754                        "IndexName": "UsernameIndex",
2755                        "KeySchema": [
2756                            {
2757                                "AttributeName": "username",
2758                                "KeyType": "HASH"
2759                            }
2760                        ],
2761                        "Projection": {
2762                            "ProjectionType": "KEYS_ONLY"
2763                        }
2764                    }
2765                ],
2766                "ProvisionedThroughput": {
2767                    "ReadCapacityUnits": 20,
2768                    "WriteCapacityUnits": 6
2769                },
2770                "TableName": "Thread",
2771                "TableStatus": "ACTIVE"
2772            }
2773        }
2774
2775        with mock.patch.object(
2776                self.users,
2777                'describe',
2778                return_value=expected) as mock_count:
2779            self.assertEqual(self.users.count(), 5)
2780
2781    def test_query_count_simple(self):
2782        expected_0 = {
2783            'Count': 0.0,
2784        }
2785
2786        expected_1 = {
2787            'Count': 10.0,
2788        }
2789
2790        with mock.patch.object(
2791                self.users.connection,
2792                'query',
2793                return_value=expected_0) as mock_query:
2794            results = self.users.query_count(username__eq='notmyname')
2795            self.assertTrue(isinstance(results, int))
2796            self.assertEqual(results, 0)
2797        self.assertEqual(mock_query.call_count, 1)
2798        self.assertIn('scan_index_forward', mock_query.call_args[1])
2799        self.assertEqual(True, mock_query.call_args[1]['scan_index_forward'])
2800        self.assertIn('limit', mock_query.call_args[1])
2801        self.assertEqual(None, mock_query.call_args[1]['limit'])
2802
2803        with mock.patch.object(
2804                self.users.connection,
2805                'query',
2806                return_value=expected_1) as mock_query:
2807            results = self.users.query_count(username__gt='somename', consistent=True, scan_index_forward=False, limit=10)
2808            self.assertTrue(isinstance(results, int))
2809            self.assertEqual(results, 10)
2810        self.assertEqual(mock_query.call_count, 1)
2811        self.assertIn('scan_index_forward', mock_query.call_args[1])
2812        self.assertEqual(False, mock_query.call_args[1]['scan_index_forward'])
2813        self.assertIn('limit', mock_query.call_args[1])
2814        self.assertEqual(10, mock_query.call_args[1]['limit'])
2815
2816    def test_query_count_paginated(self):
2817        def return_side_effect(*args, **kwargs):
2818            if kwargs.get('exclusive_start_key'):
2819                return {'Count': 10, 'LastEvaluatedKey': None}
2820            else:
2821                return {
2822                    'Count': 20,
2823                    'LastEvaluatedKey': {
2824                        'username': {
2825                            'S': 'johndoe'
2826                        },
2827                        'date_joined': {
2828                            'N': '4118642633'
2829                        }
2830                    }
2831                }
2832
2833        with mock.patch.object(
2834                self.users.connection,
2835                'query',
2836                side_effect=return_side_effect
2837        ) as mock_query:
2838            count = self.users.query_count(username__eq='johndoe')
2839            self.assertTrue(isinstance(count, int))
2840            self.assertEqual(30, count)
2841            self.assertEqual(mock_query.call_count, 2)
2842
2843    def test_private_batch_get(self):
2844        expected = {
2845            "ConsumedCapacity": {
2846                "CapacityUnits": 0.5,
2847                "TableName": "users"
2848            },
2849            'Responses': {
2850                'users': [
2851                    {
2852                        'username': {'S': 'alice'},
2853                        'first_name': {'S': 'Alice'},
2854                        'last_name': {'S': 'Expert'},
2855                        'date_joined': {'N': '1366056680'},
2856                        'friend_count': {'N': '1'},
2857                        'friends': {'SS': ['jane']},
2858                    },
2859                    {
2860                        'username': {'S': 'bob'},
2861                        'first_name': {'S': 'Bob'},
2862                        'last_name': {'S': 'Smith'},
2863                        'date_joined': {'N': '1366056888'},
2864                        'friend_count': {'N': '1'},
2865                        'friends': {'SS': ['johndoe']},
2866                    },
2867                    {
2868                        'username': {'S': 'jane'},
2869                        'first_name': {'S': 'Jane'},
2870                        'last_name': {'S': 'Doe'},
2871                        'date_joined': {'N': '1366057777'},
2872                        'friend_count': {'N': '2'},
2873                        'friends': {'SS': ['alice', 'johndoe']},
2874                    },
2875                ],
2876            },
2877            "UnprocessedKeys": {
2878            },
2879        }
2880
2881        with mock.patch.object(
2882                self.users.connection,
2883                'batch_get_item',
2884                return_value=expected) as mock_batch_get:
2885            results = self.users._batch_get(keys=[
2886                {'username': 'alice', 'friend_count': 1},
2887                {'username': 'bob', 'friend_count': 1},
2888                {'username': 'jane'},
2889            ])
2890            usernames = [res['username'] for res in results['results']]
2891            self.assertEqual(usernames, ['alice', 'bob', 'jane'])
2892            self.assertEqual(len(results['results']), 3)
2893            self.assertEqual(results['last_key'], None)
2894            self.assertEqual(results['unprocessed_keys'], [])
2895
2896        mock_batch_get.assert_called_once_with(request_items={
2897            'users': {
2898                'Keys': [
2899                    {
2900                        'username': {'S': 'alice'},
2901                        'friend_count': {'N': '1'}
2902                    },
2903                    {
2904                        'username': {'S': 'bob'},
2905                        'friend_count': {'N': '1'}
2906                    }, {
2907                        'username': {'S': 'jane'},
2908                    }
2909                ]
2910            }
2911        })
2912
2913        # Now alter the expected.
2914        del expected['Responses']['users'][2]
2915        expected['UnprocessedKeys'] = {
2916            'Keys': [
2917                {'username': {'S': 'jane',}},
2918            ],
2919        }
2920
2921        with mock.patch.object(
2922                self.users.connection,
2923                'batch_get_item',
2924                return_value=expected) as mock_batch_get_2:
2925            results = self.users._batch_get(keys=[
2926                {'username': 'alice', 'friend_count': 1},
2927                {'username': 'bob', 'friend_count': 1},
2928                {'username': 'jane'},
2929            ])
2930            usernames = [res['username'] for res in results['results']]
2931            self.assertEqual(usernames, ['alice', 'bob'])
2932            self.assertEqual(len(results['results']), 2)
2933            self.assertEqual(results['last_key'], None)
2934            self.assertEqual(results['unprocessed_keys'], [
2935                {'username': 'jane'}
2936            ])
2937
2938        mock_batch_get_2.assert_called_once_with(request_items={
2939            'users': {
2940                'Keys': [
2941                    {
2942                        'username': {'S': 'alice'},
2943                        'friend_count': {'N': '1'}
2944                    },
2945                    {
2946                        'username': {'S': 'bob'},
2947                        'friend_count': {'N': '1'}
2948                    }, {
2949                        'username': {'S': 'jane'},
2950                    }
2951                ]
2952            }
2953        })
2954
2955    def test_private_batch_get_attributes(self):
2956        # test if AttributesToGet parameter is passed to DynamoDB API
2957        expected = {
2958            "ConsumedCapacity": {
2959                "CapacityUnits": 0.5,
2960                "TableName": "users"
2961            },
2962            'Responses': {
2963                'users': [
2964                    {
2965                        'username': {'S': 'alice'},
2966                        'first_name': {'S': 'Alice'},
2967                    },
2968                    {
2969                        'username': {'S': 'bob'},
2970                        'first_name': {'S': 'Bob'},
2971                    },
2972                ],
2973            },
2974            "UnprocessedKeys": {},
2975        }
2976
2977        with mock.patch.object(
2978                self.users.connection,
2979                'batch_get_item',
2980                return_value=expected) as mock_batch_get_attr:
2981            results = self.users._batch_get(keys=[
2982                    {'username': 'alice'},
2983                    {'username': 'bob'},
2984                ], attributes=['username', 'first_name'])
2985            usernames = [res['username'] for res in results['results']]
2986            first_names = [res['first_name'] for res in results['results']]
2987            self.assertEqual(usernames, ['alice', 'bob'])
2988            self.assertEqual(first_names, ['Alice', 'Bob'])
2989            self.assertEqual(len(results['results']), 2)
2990            self.assertEqual(results['last_key'], None)
2991            self.assertEqual(results['unprocessed_keys'], [])
2992
2993        mock_batch_get_attr.assert_called_once_with(request_items={
2994            'users': {
2995                'Keys': [ { 'username': {'S': 'alice'} },
2996                          { 'username': {'S': 'bob'} }, ],
2997                'AttributesToGet': ['username', 'first_name'],
2998            },
2999        })
3000
3001    def test_batch_get(self):
3002        items_1 = {
3003            'results': [
3004                Item(self.users, data={
3005                    'username': 'johndoe',
3006                    'first_name': 'John',
3007                    'last_name': 'Doe',
3008                }),
3009                Item(self.users, data={
3010                    'username': 'jane',
3011                    'first_name': 'Jane',
3012                    'last_name': 'Doe',
3013                }),
3014            ],
3015            'last_key': None,
3016            'unprocessed_keys': [
3017                'zoeydoe',
3018            ]
3019        }
3020
3021        results = self.users.batch_get(keys=[
3022            {'username': 'johndoe'},
3023            {'username': 'jane'},
3024            {'username': 'zoeydoe'},
3025        ])
3026        self.assertTrue(isinstance(results, BatchGetResultSet))
3027        self.assertEqual(len(results._results), 0)
3028        self.assertEqual(results.the_callable, self.users._batch_get)
3029
3030        with mock.patch.object(
3031                results,
3032                'the_callable',
3033                return_value=items_1) as mock_batch_get:
3034            res_1 = next(results)
3035            # Now it should be populated.
3036            self.assertEqual(len(results._results), 2)
3037            self.assertEqual(res_1['username'], 'johndoe')
3038            res_2 = next(results)
3039            self.assertEqual(res_2['username'], 'jane')
3040
3041        self.assertEqual(mock_batch_get.call_count, 1)
3042        self.assertEqual(results._keys_left, ['zoeydoe'])
3043
3044        items_2 = {
3045            'results': [
3046                Item(self.users, data={
3047                    'username': 'zoeydoe',
3048                    'first_name': 'Zoey',
3049                    'last_name': 'Doe',
3050                }),
3051            ],
3052        }
3053
3054        with mock.patch.object(
3055                results,
3056                'the_callable',
3057                return_value=items_2) as mock_batch_get_2:
3058            res_3 = next(results)
3059            # New results should have been found.
3060            self.assertEqual(len(results._results), 1)
3061            self.assertEqual(res_3['username'], 'zoeydoe')
3062
3063            self.assertRaises(StopIteration, results.next)
3064
3065        self.assertEqual(mock_batch_get_2.call_count, 1)
3066        self.assertEqual(results._keys_left, [])
3067