def _call_test_fn(self, table_name='my-table'): table = Table(table_name) put_attributes = {'foo': 'bar'} update_arg = db.UpdateArg(self._pk, self._sk, attr_updates=put_attributes) return table._update_item(update_arg)
def setUp(self): super().setUp() self._pk_2 = db.PartitionKey(User, '*****@*****.**') self._sk_2 = db.SortKey(Subscription, 'docs.bar.com') self._keys = [ db.PrimaryKey(self._pk, self._sk), db.PrimaryKey(self._pk_2, self._sk_2) ] self._table_name = 'my-table' self._table = Table(self._table_name)
def test_primary_index(self): pk_name = 'my-pk-name' sk_name = 'my-sk-name' class TestIndex(db.GlobalIndex): @property def partition_key(self) -> str: return pk_name @property def sort_key(self) -> str: return sk_name table = Table('my-table', TestIndex()) self.assertEqual(table.primary_index.partition_key, pk_name) self.assertEqual(table.primary_index.sort_key, sk_name)
def test_client(self): boto3 = self._mocks['boto3'] table = Table('my-table') self.assertEqual(table._client, boto3.client.return_value)
def test_makes_copy(self): item = {'PK': str(self._pk), 'SK': str(self._sk)} res = Table._strip_prefixes(item) self.assertNotEqual(item['PK'], res['PK']) self.assertNotEqual(item['SK'], res['SK'])
def test_strips_prefixes(self): item = {'PK': str(self._pk), 'SK': str(self._sk), 'Foo': str(self._sk)} res = Table._strip_prefixes(item) self.assertEqual(res['PK'], self._pk.value) self.assertEqual(res['SK'], self._sk.value) self.assertEqual(res['Foo'], self._sk.value)
def test_noop_on_no_prefix(self): item = {'foo': 'bar'} res = Table._strip_prefixes(item) self.assertDictEqual(item, res)
def _call_test_fn(self, items=None, table_name='my-table'): table = Table(table_name) if not items: items = [] return table.transact_write_items(items)
def _call_test_fn(self, table_name='my-table'): table = Table(table_name) put_arg = db.PutArg(self._pk, self._sk) return table._put_item(put_arg)
def _call_test_fn(self, global_index=None, attributes=None): table = Table('my-table') return table.query_prefix(self._pk, self._sk_prefix, global_index=global_index, attributes=attributes)
def test_noop_on_no_match(self): val = 'foo' res = Table._remove_entity_prefix(val) self.assertEqual(val, res)
def _call_test_fn(self, attributes=None): table = Table('my-table') return table.get(self._pk, self._sk, attributes=attributes)
def _call_test_fn(self, table_name='my-table'): table = Table(table_name) key_cond = Key('PK').eq(str(self._pk)) query_arg = db.QueryArg(key_cond) return table._query(query_arg)
def _test_val(self, prefix, val): res = Table._remove_entity_prefix(f'{prefix}{val}') self.assertEqual(val, res)
def _call_test_fn(self, table_name='my-table'): table = Table(table_name) return table.delete(self._pk, self._sk)
class TestBatchGet(TableTestCaseMixin, TestBase): def setUp(self): super().setUp() self._pk_2 = db.PartitionKey(User, '*****@*****.**') self._sk_2 = db.SortKey(Subscription, 'docs.bar.com') self._keys = [ db.PrimaryKey(self._pk, self._sk), db.PrimaryKey(self._pk_2, self._sk_2) ] self._table_name = 'my-table' self._table = Table(self._table_name) def _call_test_fn(self, attributes=None, consistent=False): attributes = attributes or [] return self._table.batch_get(self._keys, attributes=attributes, consistent=consistent) def _get_call_arg(self, name, consistent=False, attributes=None): self._call_test_fn(attributes=attributes, consistent=consistent) _, kwargs = self._dynamo_method.call_args return kwargs['RequestItems'][self._table_name][name] def _get_attributes_call_arg(self, attributes=None): pe = self._get_call_arg('ProjectionExpression', attributes) attributes = pe.split(',') return attributes @property def _dynamo_method(self): return self._client.batch_get_item def test_keys(self): keys = self._get_call_arg('Keys', consistent=True) self.assertEqual(len(keys), 2) for i, key in enumerate(keys): exp = self._keys[i].serialize(self._table.primary_index) self.assertDictEqual(keys[i], exp) def test_consistent(self): consistent = self._get_call_arg('ConsistentRead', consistent=True) self.assertTrue(consistent) def test_retrieves_keys(self): attributes = self._get_attributes_call_arg(['foo']) self.assertTrue(set(attributes).issuperset({'PK', 'SK'})) def test_retrieves_keys_default(self): attributes = self._get_attributes_call_arg() self.assertSetEqual(set(attributes), {'PK', 'SK'}) def test_results(self): key_1_ser = self._keys[0].serialize(self._table.primary_index) key_2_ser = self._keys[1].serialize(self._table.primary_index) self._dynamo_method.return_value = { 'Responses': { self._table_name: [key_1_ser] }, 'UnprocessedKeys': { self._table_name: { 'Keys': [key_2_ser] } } } res = self._call_test_fn() self.assertEqual(len(res.items), 1) self.assertEqual(res.items[0]['PK'], self._pk.value) self.assertEqual(res.items[0]['SK'], self._sk.value) self.assertEqual(len(res.unprocessed_keys), 1) self.assertEqual(res.unprocessed_keys[0], self._keys[1])