117 lines
4.3 KiB
Python
117 lines
4.3 KiB
Python
# Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License"). You
|
|
# may not use this file except in compliance with the License. A copy of
|
|
# the License is located at
|
|
#
|
|
# https://aws.amazon.com/apache2.0/
|
|
#
|
|
# or in the "license" file accompanying this file. This file is
|
|
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
|
|
# ANY KIND, either express or implied. See the License for the specific
|
|
# language governing permissions and limitations under the License.
|
|
import json
|
|
|
|
from botocore.stub import Stubber
|
|
|
|
import boto3
|
|
from tests import unittest
|
|
|
|
|
|
class TestTableResourceCustomizations(unittest.TestCase):
|
|
|
|
maxDiff = None
|
|
|
|
def setUp(self):
|
|
self.resource = boto3.resource('dynamodb', 'us-east-1')
|
|
self.table = self.resource.Table('mytable')
|
|
|
|
def test_resource_has_batch_writer_added(self):
|
|
assert hasattr(self.table, 'batch_writer')
|
|
|
|
def test_operation_without_output(self):
|
|
stubber = Stubber(self.table.meta.client)
|
|
stubber.add_response('tag_resource', {})
|
|
arn = 'arn:aws:dynamodb:us-west-2:123456789:table/mytable'
|
|
|
|
with stubber:
|
|
self.table.meta.client.tag_resource(
|
|
ResourceArn=arn, Tags=[{'Key': 'project', 'Value': 'val'}]
|
|
)
|
|
|
|
stubber.assert_no_pending_responses()
|
|
|
|
def test_batch_write_does_not_double_serialize(self):
|
|
# If multiple items reference the same Python object, the
|
|
# object does not get double-serialized.
|
|
# https://github.com/boto/boto3/issues/3474
|
|
|
|
used_twice = {'pkey': 'foo1', 'otherfield': {'foo': 1, 'bar': 2}}
|
|
batch_writer = self.table.batch_writer()
|
|
|
|
# The default Stubber compares the request payload to the
|
|
# "expected_params" before automatic serialization happens. This custom
|
|
# event handler uses the same technique as the Stubber to record the
|
|
# serialized request body, but later in the request lifecycle.
|
|
class LateStubber:
|
|
def __init__(self, client):
|
|
self.intercepted_request_body = None
|
|
client.meta.events.register_first(
|
|
'before-call.*.*',
|
|
self.late_request_interceptor,
|
|
)
|
|
|
|
def late_request_interceptor(self, event_name, params, **kwargs):
|
|
if self.intercepted_request_body is not None:
|
|
raise AssertionError(
|
|
'LateStubber was called more than once, but only one '
|
|
'request is expected'
|
|
)
|
|
body_str = params.get('body', b'').decode('utf-8')
|
|
try:
|
|
self.intercepted_request_body = json.loads(body_str)
|
|
except Exception:
|
|
raise AssertionError(
|
|
'Expected JSON request body, but failed to JSON decode'
|
|
)
|
|
|
|
late_stubber = LateStubber(self.table.meta.client)
|
|
|
|
with Stubber(self.table.meta.client) as stubber:
|
|
stubber.add_response(
|
|
'batch_write_item',
|
|
service_response={'UnprocessedItems': {}},
|
|
)
|
|
batch_writer.put_item(Item=used_twice)
|
|
batch_writer.put_item(Item=used_twice)
|
|
batch_writer._flush()
|
|
|
|
expected_request_body = {
|
|
'RequestItems': {
|
|
'mytable': [
|
|
{
|
|
'PutRequest': {
|
|
'Item': {
|
|
'pkey': {'S': 'foo1'},
|
|
'otherfield': {
|
|
'M': {'foo': {'N': '1'}, 'bar': {'N': '2'}}
|
|
},
|
|
}
|
|
}
|
|
},
|
|
{
|
|
'PutRequest': {
|
|
'Item': {
|
|
'pkey': {'S': 'foo1'},
|
|
'otherfield': {
|
|
'M': {'foo': {'N': '1'}, 'bar': {'N': '2'}}
|
|
},
|
|
}
|
|
}
|
|
},
|
|
]
|
|
}
|
|
}
|
|
|
|
assert late_stubber.intercepted_request_body == expected_request_body
|