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])
