AioStubber that returns AioAWSResponse()

aiobotocore/stub.py

@@ -0,0 +1,107 @@
+from .awsrequest import AioAWSResponse
+from botocore.stub import Stubber
+
+
+class AioStubber(Stubber):
+	def _add_response(self, method, service_response, expected_params):
+		if not hasattr(self.client, method):
+			raise ValueError(
+				"Client %s does not have method: %s"
+				% (self.client.meta.service_model.service_name, method)
+			)
+
+		# Create a successful http response
+		http_response = AioAWSResponse(None, 200, {}, None)
+
+		operation_name = self.client.meta.method_to_api_mapping.get(method)
+		self._validate_operation_response(operation_name, service_response)
+
+		# Add the service_response to the queue for returning responses
+		response = {
+			'operation_name': operation_name,
+			'response': (http_response, service_response),
+			'expected_params': expected_params,
+		}
+		self._queue.append(response)
+
+	def add_client_error(
+			self,
+			method,
+			service_error_code='',
+			service_message='',
+			http_status_code=400,
+			service_error_meta=None,
+			expected_params=None,
+			response_meta=None,
+			modeled_fields=None,
+	):
+		"""
+		Adds a ``ClientError`` to the response queue.
+
+		:param method: The name of the service method to return the error on.
+		:type method: str
+
+		:param service_error_code: The service error code to return,
+								   e.g. ``NoSuchBucket``
+		:type service_error_code: str
+
+		:param service_message: The service message to return, e.g.
+						'The specified bucket does not exist.'
+		:type service_message: str
+
+		:param http_status_code: The HTTP status code to return, e.g. 404, etc
+		:type http_status_code: int
+
+		:param service_error_meta: Additional keys to be added to the
+			service Error
+		:type service_error_meta: dict
+
+		:param expected_params: A dictionary of the expected parameters to
+			be called for the provided service response. The parameters match
+			the names of keyword arguments passed to that client call. If
+			any of the parameters differ a ``StubResponseError`` is thrown.
+			You can use stub.ANY to indicate a particular parameter to ignore
+			in validation.
+
+		:param response_meta: Additional keys to be added to the
+			response's ResponseMetadata
+		:type response_meta: dict
+
+		:param modeled_fields: Additional keys to be added to the response
+			based on fields that are modeled for the particular error code.
+			These keys will be validated against the particular error shape
+			designated by the error code.
+		:type modeled_fields: dict
+
+		"""
+		http_response = AioAWSResponse(None, http_status_code, {}, None)
+
+		# We don't look to the model to build this because the caller would
+		# need to know the details of what the HTTP body would need to
+		# look like.
+		parsed_response = {
+			'ResponseMetadata': {'HTTPStatusCode': http_status_code},
+			'Error': {'Message': service_message, 'Code': service_error_code},
+		}
+
+		if service_error_meta is not None:
+			parsed_response['Error'].update(service_error_meta)
+
+		if response_meta is not None:
+			parsed_response['ResponseMetadata'].update(response_meta)
+
+		if modeled_fields is not None:
+			service_model = self.client.meta.service_model
+			shape = service_model.shape_for_error_code(service_error_code)
+			self._validate_response(shape, modeled_fields)
+			parsed_response.update(modeled_fields)
+
+		operation_name = self.client.meta.method_to_api_mapping.get(method)
+		# Note that we do not allow for expected_params while
+		# adding errors into the queue yet.
+		response = {
+			'operation_name': operation_name,
+			'response': (http_response, parsed_response),
+			'expected_params': expected_params,
+		}
+		self._queue.append(response)