def test_consumed_capacity_equality(self):
""" ConsumedCapacity addition and equality """
cap = ConsumedCapacity('foobar', Capacity(0, 10), Capacity(0, 2), {
'l-index': Capacity(0, 4),
}, {
'g-index': Capacity(0, 3),
})
c2 = ConsumedCapacity('foobar', Capacity(0, 10), Capacity(0, 2), {
'l-index': Capacity(0, 4),
'l-index2': Capacity(0, 7),
})
self.assertNotEqual(cap, c2)
c3 = ConsumedCapacity('foobar', Capacity(0, 10), Capacity(0, 2), {
'l-index': Capacity(0, 4),
}, {
'g-index': Capacity(0, 3),
})
self.assertIn(cap, set([c3]))
combined = cap + c2
self.assertEqual(
cap + c2,
ConsumedCapacity('foobar', Capacity(0, 20), Capacity(0, 4), {
'l-index': Capacity(0, 8),
'l-index2': Capacity(0, 7),
}, {
'g-index': Capacity(0, 3),
}))
self.assertIn(str(Capacity(0, 3)), str(combined))
def test_consumed_capacity_equality(self):
"""ConsumedCapacity addition and equality"""
cap = ConsumedCapacity(
"foobar",
Capacity(0, 10),
Capacity(0, 2),
{
"l-index": Capacity(0, 4),
},
{
"g-index": Capacity(0, 3),
},
)
c2 = ConsumedCapacity(
"foobar",
Capacity(0, 10),
Capacity(0, 2),
{
"l-index": Capacity(0, 4),
"l-index2": Capacity(0, 7),
},
)
self.assertNotEqual(cap, c2)
c3 = ConsumedCapacity(
"foobar",
Capacity(0, 10),
Capacity(0, 2),
{
"l-index": Capacity(0, 4),
},
{
"g-index": Capacity(0, 3),
},
)
self.assertIn(cap, set([c3]))
combined = cap + c2
self.assertEqual(
cap + c2,
ConsumedCapacity(
"foobar",
Capacity(0, 20),
Capacity(0, 4),
{
"l-index": Capacity(0, 8),
"l-index2": Capacity(0, 7),
},
{
"g-index": Capacity(0, 3),
},
),
)
self.assertIn(str(Capacity(0, 3)), str(combined))
def test_throttle_table_default_cap(self):
""" If no table limit provided, use the default """
limiter = RateLimit(default=Capacity(4, 4))
cap = ConsumedCapacity('foobar', Capacity(8, 0), Capacity(8, 0))
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query2('foobar', 'id = :id', id='a'))
sleep.assert_called_with(2)
def test_no_throttle(self):
"""Don't sleep if consumed capacity is within limits"""
limiter = RateLimit(3, 3)
cap = ConsumedCapacity("foobar", Capacity(0, 2))
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_not_called()
def test_throttle_total_cap(self):
"""Sleep if consumed capacity exceeds total"""
limiter = RateLimit(total=Capacity(3, 3))
cap = ConsumedCapacity("foobar", Capacity(3, 0))
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_called_with(1)
def test_throttle_multiply(self):
"""Seconds to sleep is increades to match limit delta"""
limiter = RateLimit(3, 3)
cap = ConsumedCapacity("foobar", Capacity(8, 0))
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_called_with(3)
def test_throttle_total(self):
""" Sleep if consumed capacity exceeds total """
limiter = RateLimit(3, 3)
cap = ConsumedCapacity('foobar', Capacity(3, 0))
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query2('foobar', 'id = :id', id='a'))
sleep.assert_called_with(1)
def test_throttle_table_default(self):
"""If no table limit provided, use the default"""
limiter = RateLimit(default_read=4, default_write=4)
cap = ConsumedCapacity("foobar", Capacity(8, 0), Capacity(8, 0))
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_called_with(2)
def test_throttle_multiple(self):
"""Sleep if the limit is exceeded by multiple calls"""
limiter = RateLimit(4, 4)
cap = ConsumedCapacity("foobar", Capacity(3, 0))
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_called_with(2)
def test_total_consumed_capacity(self):
""" ConsumedCapacity can parse results with only Total """
response = {
'TableName': 'foobar',
'CapacityUnits': 4,
}
cap = ConsumedCapacity.from_response(response, True)
self.assertEqual(cap.total.read, 4)
self.assertIsNone(cap.table_capacity)
def test_throttle_table(self):
""" Sleep if table limit is exceeded """
limiter = RateLimit(3, 3, table_caps={
'foobar': Capacity(0, 4),
})
cap = ConsumedCapacity('foobar', Capacity(8, 0), Capacity(0, 8))
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query2('foobar', 'id = :id', id='a'))
sleep.assert_called_with(2)
def test_total_consumed_capacity(self):
""" ConsumedCapacity can parse results with only Total """
response = {
'TableName': 'foobar',
'CapacityUnits': 4,
}
cap = ConsumedCapacity.from_response(response, True)
self.assertEqual(cap.total.read, 4)
self.assertIsNone(cap.table_capacity)
def test_throttle_callback(self):
"""Callback is called when a query is throttled"""
callback = MagicMock()
callback.return_value = True
limiter = RateLimit(3, 3, callback=callback)
cap = ConsumedCapacity("foobar", Capacity(3, 0))
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_not_called()
self.assertTrue(callback.called)
def test_total_consumed_capacity(self):
"""ConsumedCapacity can parse results with only Total"""
response = {
"TableName": "foobar",
"ReadCapacityUnits": 4,
"WriteCapacityUnits": 5,
}
cap = ConsumedCapacity.from_response(response)
self.assertEqual(cap.total, (4, 5))
self.assertIsNone(cap.table_capacity)
def test_global_default(self):
""" Global index limit will fall back to table default limit """
limiter = RateLimit(default_read=4, default_write=4)
cap = ConsumedCapacity('foobar',
Capacity(8, 0),
global_index_capacity={
'baz': Capacity(8, 0),
})
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query2('foobar', 'id = :id', id='a'))
sleep.assert_called_with(2)
def test_global_default_table(self):
""" Global index limit defaults to table limit if not present """
limiter = RateLimit(table_caps={
'foobar': Capacity(4, 0),
})
cap = ConsumedCapacity('foobar',
Capacity(8, 0),
global_index_capacity={
'baz': Capacity(8, 0),
})
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query2('foobar', 'id = :id', id='a'))
sleep.assert_called_with(2)
def test_throttle_table(self):
"""Sleep if table limit is exceeded"""
limiter = RateLimit(
3,
3,
table_caps={
"foobar": Capacity(0, 4),
},
)
cap = ConsumedCapacity("foobar", Capacity(8, 0), Capacity(0, 8))
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_called_with(2)
def test_local_index(self):
""" Local index capacities count towards the table limit """
limiter = RateLimit(table_caps={
'foobar': Capacity(4, 0),
})
cap = ConsumedCapacity('foobar',
Capacity(8, 0),
local_index_capacity={
'local': Capacity(4, 0),
})
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query2('foobar', 'id = :id', id='a'))
sleep.assert_called_with(1)
def test_global_index(self):
""" Sleep when global index limit is exceeded """
limiter = RateLimit(table_caps={'foobar': {
'baz': Capacity(4, 0),
}})
cap = ConsumedCapacity('foobar',
Capacity(8, 0),
global_index_capacity={
'baz': Capacity(8, 0),
})
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query2('foobar', 'id = :id', id='a'))
sleep.assert_called_with(2)
def test_global_default(self):
"""Global index limit will fall back to table default limit"""
limiter = RateLimit(default_read=4, default_write=4)
cap = ConsumedCapacity(
"foobar",
Capacity(8, 0),
global_index_capacity={
"baz": Capacity(8, 0),
},
)
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_called_with(2)
def test_global_index_by_name(self):
""" Global index limit can be specified as tablename:index_name """
limiter = RateLimit(table_caps={
'foobar:baz': Capacity(4, 0),
})
cap = ConsumedCapacity('foobar',
Capacity(8, 0),
global_index_capacity={
'baz': Capacity(8, 0),
})
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query2('foobar', 'id = :id', id='a'))
sleep.assert_called_with(2)
def test_global_index(self):
"""Sleep when global index limit is exceeded"""
limiter = RateLimit(table_caps={"foobar": {
"baz": Capacity(4, 0),
}})
cap = ConsumedCapacity(
"foobar",
Capacity(8, 0),
global_index_capacity={
"baz": Capacity(8, 0),
},
)
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_called_with(2)
def test_global_default_table(self):
"""Global index limit defaults to table limit if not present"""
limiter = RateLimit(table_caps={
"foobar": Capacity(4, 0),
})
cap = ConsumedCapacity(
"foobar",
Capacity(8, 0),
global_index_capacity={
"baz": Capacity(8, 0),
},
)
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_called_with(2)
def test_global_index_by_name(self):
"""Global index limit can be specified as tablename:index_name"""
limiter = RateLimit(table_caps={
"foobar:baz": Capacity(4, 0),
})
cap = ConsumedCapacity(
"foobar",
Capacity(8, 0),
global_index_capacity={
"baz": Capacity(8, 0),
},
)
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_called_with(2)
def test_local_index(self):
"""Local index capacities count towards the table limit"""
limiter = RateLimit(table_caps={
"foobar": Capacity(4, 0),
})
cap = ConsumedCapacity(
"foobar",
Capacity(8, 0),
local_index_capacity={
"local": Capacity(4, 0),
},
)
with self.inject_capacity(cap, limiter) as sleep:
list(self.dynamo.query("foobar", "id = :id", id="a"))
sleep.assert_called_with(1)
def test_capacity(self):
"""Can return consumed capacity"""
conn = MagicMock()
response_cap = {
"TableName": "foobar",
"ReadCapacityUnits": 6,
"WriteCapacityUnits": 7,
"Table": {
"ReadCapacityUnits": 1,
"WriteCapacityUnits": 2,
},
"LocalSecondaryIndexes": {
"l-index": {
"ReadCapacityUnits": 2,
"WriteCapacityUnits": 3,
},
},
"GlobalSecondaryIndexes": {
"g-index": {
"ReadCapacityUnits": 3,
"WriteCapacityUnits": 4,
},
},
}
response = {
"Responses": [],
"ConsumedCapacity": [response_cap],
}
capacity = ConsumedCapacity.from_response(response_cap)
response["consumed_capacity"] = [capacity]
with patch.object(self.dynamo, "client") as client:
client.transact_get_items.return_value = response
ret = self.dynamo.txn_get("foobar", [{"id": "a"}])
list(ret)
assert ret.consumed_capacity is not None
cap = ret.consumed_capacity["foobar"]
assert cap is not None
assert cap.table_capacity is not None
assert cap.local_index_capacity is not None
assert cap.global_index_capacity is not None
self.assertEqual(cap.total, Throughput(6, 7))
self.assertEqual(cap.table_capacity, Throughput(1, 2))
self.assertEqual(cap.local_index_capacity["l-index"], Throughput(2, 3))
self.assertEqual(cap.global_index_capacity["g-index"],
Throughput(3, 4))
def test_capacity(self):
"""capacity hooks are called whenever response has ConsumedCapacity"""
hash_key = DynamoKey("id")
self.dynamo.create_table("foobar", hash_key=hash_key)
hook = MagicMock()
self.dynamo.subscribe("capacity", hook)
with patch.object(self.dynamo, "client") as client:
client.scan.return_value = {
"Items": [],
"ConsumedCapacity": {
"TableName": "foobar",
"ReadCapacityUnits": 4,
},
}
rs = self.dynamo.scan("foobar")
list(rs)
cap = ConsumedCapacity("foobar", Capacity(4, 0))
hook.assert_called_with(self.dynamo, "scan", ANY, ANY, cap)
def test_capacity(self):
""" capacity hooks are called whenever response has ConsumedCapacity """
hash_key = DynamoKey('id')
self.dynamo.create_table('foobar', hash_key=hash_key)
hook = MagicMock()
self.dynamo.subscribe('capacity', hook)
with patch.object(self.dynamo, 'client') as client:
client.scan.return_value = {
'Items': [],
'ConsumedCapacity': {
'TableName': 'foobar',
'CapacityUnits': 4,
}
}
rs = self.dynamo.scan('foobar')
list(rs)
cap = ConsumedCapacity('foobar', Capacity(4, 0))
hook.assert_called_with(self.dynamo, 'scan', ANY, ANY, cap)
def test_capacity(self):
"""Can return consumed capacity"""
conn = MagicMock()
response_cap = {
"TableName": "foobar",
"CapacityUnits": 6,
"Table": {
"CapacityUnits": 1,
},
"LocalSecondaryIndexes": {
"l-index": {
"CapacityUnits": 2,
},
},
"GlobalSecondaryIndexes": {
"g-index": {
"CapacityUnits": 3,
},
},
}
response = {
"Responses": {
"foo": [],
},
"ConsumedCapacity": [response_cap],
}
capacity = ConsumedCapacity.from_response(response_cap, True)
response["consumed_capacity"] = [capacity]
conn.call.return_value = response
rs = GetResultSet(conn, {"foobar": [{"id": "a"}]})
list(rs)
assert rs.consumed_capacity is not None
cap = rs.consumed_capacity["foobar"]
assert cap is not None
assert cap.table_capacity is not None
assert cap.local_index_capacity is not None
assert cap.global_index_capacity is not None
self.assertEqual(cap.total, (6, 0))
self.assertEqual(cap.table_capacity, (1, 0))
self.assertEqual(cap.local_index_capacity["l-index"], (2, 0))
self.assertEqual(cap.global_index_capacity["g-index"], (3, 0))
def test_capacity(self):
""" Can return consumed capacity """
conn = MagicMock()
response = {
'Responses': {
'foo': [],
},
'ConsumedCapacity': [{
'TableName': 'foobar',
'CapacityUnits': 3,
'Table': {
'CapacityUnits': 1,
},
'LocalSecondaryIndexes': {
'l-index': {
'CapacityUnits': 1,
},
},
'GlobalSecondaryIndexes': {
'g-index': {
'CapacityUnits': 1,
},
},
}],
}
capacity = ConsumedCapacity.from_response(
response['ConsumedCapacity'][0], True)
response['consumed_capacity'] = [capacity]
conn.call.return_value = response
rs = GetResultSet(conn, 'foo', [{'id': 'a'}])
list(rs)
self.assertEqual(rs.capacity, 3)
self.assertEqual(rs.table_capacity, 1)
self.assertEqual(rs.indexes, {'l-index': 1})
self.assertEqual(rs.global_indexes, {'g-index': 1})
self.assertEqual(rs.consumed_capacity, capacity)
def test_add_different_tables(self):
""" Cannot add ConsumedCapacity of two different tables """
c1 = ConsumedCapacity('foobar', Capacity(1, 28))
c2 = ConsumedCapacity('boofar', Capacity(3, 0))
with self.assertRaises(TypeError):
c1 += c2
