├── .coveragerc ├── .github └── PULL_REQUEST_TEMPLATE.md ├── .gitignore ├── .travis.yml ├── CODE_OF_CONDUCT.md ├── CONTRIBUTING.md ├── LICENSE ├── NOTICE ├── README.md ├── crhelper ├── __init__.py ├── __init__.pyi ├── log_helper.py ├── log_helper.pyi ├── py.typed ├── resource_helper.py ├── resource_helper.pyi ├── utils.py └── utils.pyi ├── requirements.txt ├── setup.py └── tests ├── __init__.py ├── functional ├── create.json └── lambda_function.py ├── test_log_helper.py ├── test_resource_helper.py ├── test_utils.py └── unit └── __init__.py /.coveragerc: -------------------------------------------------------------------------------- 1 | [run] 2 | branch = True 3 | include = crhelper/* -------------------------------------------------------------------------------- /.github/PULL_REQUEST_TEMPLATE.md: -------------------------------------------------------------------------------- 1 | *Issue #, if available:* 2 | 3 | *Description of changes:* 4 | 5 | 6 | By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license. 7 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /.idea/ 2 | /venv/ 3 | /.eggs/ 4 | __pycache__ 5 | *.pyc 6 | /dist/ 7 | .coverage 8 | coverage.xml 9 | /.pytest_cache/ 10 | /build/ 11 | *.egg-info/ 12 | /.mypy_cache/ 13 | .env/ -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | dist: xenial 2 | language: python 3 | python: 4 | - "3.6" 5 | - "3.7" 6 | install: 7 | - pip install -r requirements.txt 8 | - pip install . 9 | - pip install pytest-cov codecov 10 | script: pytest --cov=./ 11 | after_success: 12 | - codecov --token=$CODECOV_TOKEN 13 | -------------------------------------------------------------------------------- /CODE_OF_CONDUCT.md: -------------------------------------------------------------------------------- 1 | ## Code of Conduct 2 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 3 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 4 | opensource-codeofconduct@amazon.com with any additional questions or comments. 5 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | # Contributing Guidelines 2 | 3 | Thank you for your interest in contributing to our project. Whether it's a bug report, new feature, correction, or additional 4 | documentation, we greatly value feedback and contributions from our community. 5 | 6 | Please read through this document before submitting any issues or pull requests to ensure we have all the necessary 7 | information to effectively respond to your bug report or contribution. 8 | 9 | 10 | ## Reporting Bugs/Feature Requests 11 | 12 | We welcome you to use the GitHub issue tracker to report bugs or suggest features. 13 | 14 | When filing an issue, please check [existing open](https://github.com/aws-cloudformation/custom-resource-helper/issues), or [recently closed](https://github.com/aws-cloudformation/custom-resource-helper/issues?utf8=%E2%9C%93&q=is%3Aissue%20is%3Aclosed%20), issues to make sure somebody else hasn't already 15 | reported the issue. Please try to include as much information as you can. Details like these are incredibly useful: 16 | 17 | * A reproducible test case or series of steps 18 | * The version of our code being used 19 | * Any modifications you've made relevant to the bug 20 | * Anything unusual about your environment or deployment 21 | 22 | 23 | ## Contributing via Pull Requests 24 | Contributions via pull requests are much appreciated. Before sending us a pull request, please ensure that: 25 | 26 | 1. You are working against the latest source on the *main* branch. 27 | 2. You check existing open, and recently merged, pull requests to make sure someone else hasn't addressed the problem already. 28 | 3. You open an issue to discuss any significant work - we would hate for your time to be wasted. 29 | 30 | To send us a pull request, please: 31 | 32 | 1. Fork the repository. 33 | 2. Modify the source; please focus on the specific change you are contributing. If you also reformat all the code, it will be hard for us to focus on your change. 34 | 3. Ensure local tests pass. 35 | 4. Commit to your fork using clear commit messages. 36 | 5. Send us a pull request, answering any default questions in the pull request interface. 37 | 6. Pay attention to any automated CI failures reported in the pull request, and stay involved in the conversation. 38 | 39 | GitHub provides additional document on [forking a repository](https://help.github.com/articles/fork-a-repo/) and 40 | [creating a pull request](https://help.github.com/articles/creating-a-pull-request/). 41 | 42 | 43 | ## Finding contributions to work on 44 | Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any ['help wanted'](https://github.com/aws-cloudformation/custom-resource-helper/labels/help%20wanted) issues is a great place to start. 45 | 46 | 47 | ## Code of Conduct 48 | This project has adopted the [Amazon Open Source Code of Conduct](https://aws.github.io/code-of-conduct). 49 | For more information see the [Code of Conduct FAQ](https://aws.github.io/code-of-conduct-faq) or contact 50 | opensource-codeofconduct@amazon.com with any additional questions or comments. 51 | 52 | 53 | ## Security issue notifications 54 | If you discover a potential security issue in this project we ask that you notify AWS/Amazon Security via our [vulnerability reporting page](http://aws.amazon.com/security/vulnerability-reporting/). Please do **not** create a public github issue. 55 | 56 | 57 | ## Licensing 58 | 59 | See the [LICENSE](https://github.com/aws-cloudformation/custom-resource-helper/blob/main/LICENSE) file for our project's licensing. We will ask you to confirm the licensing of your contribution. 60 | 61 | We may ask you to sign a [Contributor License Agreement (CLA)](http://en.wikipedia.org/wiki/Contributor_License_Agreement) for larger changes. 62 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | -------------------------------------------------------------------------------- /NOTICE: -------------------------------------------------------------------------------- 1 | Custom Resource Helper 2 | Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved. 3 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## Custom Resource Helper 2 | 3 | Simplify best practice Custom Resource creation, sending responses to CloudFormation and providing exception, timeout 4 | trapping, and detailed configurable logging. 5 | 6 | [![PyPI Version](https://img.shields.io/pypi/v/crhelper.svg)](https://pypi.org/project/crhelper/) 7 | ![Python Versions](https://img.shields.io/pypi/pyversions/crhelper.svg) 8 | [![Build Status](https://travis-ci.com/aws-cloudformation/custom-resource-helper.svg?branch=main)](https://travis-ci.com/aws-cloudformation/custom-resource-helper) 9 | [![Test Coverage](https://codecov.io/gh/aws-cloudformation/custom-resource-helper/branch/main/graph/badge.svg)](https://codecov.io/gh/aws-cloudformation/custom-resource-helper) 10 | 11 | ## Features 12 | 13 | * Dead simple to use, reduces the complexity of writing a CloudFormation custom resource 14 | * Guarantees that CloudFormation will get a response even if an exception is raised 15 | * Returns meaningful errors to CloudFormation Stack events in the case of a failure 16 | * Polling enables run times longer than the lambda 15 minute limit 17 | * JSON logging that includes request id's, stack id's and request type to assist in tracing logs relevant to a 18 | particular CloudFormation event 19 | * Catches function timeouts and sends CloudFormation a failure response 20 | * Static typing (mypy) compatible 21 | 22 | ## Installation 23 | 24 | Install into the root folder of your lambda function 25 | 26 | ```shell 27 | cd my-lambda-function/ 28 | pip install crhelper -t . 29 | ``` 30 | 31 | ## Example Usage 32 | 33 | [This blog](https://aws.amazon.com/blogs/infrastructure-and-automation/aws-cloudformation-custom-resource-creation-with-python-aws-lambda-and-crhelper/) covers usage in more detail. 34 | 35 | ```python 36 | from __future__ import print_function 37 | from crhelper import CfnResource 38 | import logging 39 | 40 | logger = logging.getLogger(__name__) 41 | # Initialise the helper, all inputs are optional, this example shows the defaults 42 | helper = CfnResource(json_logging=False, log_level='DEBUG', boto_level='CRITICAL', sleep_on_delete=120, ssl_verify=None) 43 | 44 | try: 45 | ## Init code goes here 46 | pass 47 | except Exception as e: 48 | helper.init_failure(e) 49 | 50 | 51 | @helper.create 52 | def create(event, context): 53 | logger.info("Got Create") 54 | # Optionally return an ID that will be used for the resource PhysicalResourceId, 55 | # if None is returned an ID will be generated. If a poll_create function is defined 56 | # return value is placed into the poll event as event['CrHelperData']['PhysicalResourceId'] 57 | # 58 | # To add response data update the helper.Data dict 59 | # If poll is enabled data is placed into poll event as event['CrHelperData'] 60 | helper.Data.update({"test": "testdata"}) 61 | 62 | # To return an error to cloudformation you raise an exception: 63 | if not helper.Data.get("test"): 64 | raise ValueError("this error will show in the cloudformation events log and console.") 65 | 66 | return "MyResourceId" 67 | 68 | 69 | @helper.update 70 | def update(event, context): 71 | logger.info("Got Update") 72 | # If the update resulted in a new resource being created, return an id for the new resource. 73 | # CloudFormation will send a delete event with the old id when stack update completes 74 | 75 | 76 | @helper.delete 77 | def delete(event, context): 78 | logger.info("Got Delete") 79 | # Delete never returns anything. Should not fail if the underlying resources are already deleted. 80 | # Desired state. 81 | 82 | 83 | @helper.poll_create 84 | def poll_create(event, context): 85 | logger.info("Got create poll") 86 | # Return a resource id or True to indicate that creation is complete. if True is returned an id 87 | # will be generated 88 | return True 89 | 90 | 91 | def handler(event, context): 92 | helper(event, context) 93 | ``` 94 | 95 | ### Polling 96 | 97 | If you need longer than the max runtime of 15 minutes, you can enable polling by adding additional decorators for 98 | `poll_create`, `poll_update` or `poll_delete`. When a poll function is defined for `create`/`update`/`delete` the 99 | function will not send a response to CloudFormation and instead a CloudWatch Events schedule will be created to 100 | re-invoke the lambda function every 2 minutes. When the function is invoked the matching `@helper.poll_` function will 101 | be called, logic to check for completion should go here, if the function returns `None` then the schedule will run again 102 | in 2 minutes. Once complete either return a PhysicalResourceID or `True` to have one generated. The schedule will be 103 | deleted and a response sent back to CloudFormation. If you use polling the following additional IAM policy must be 104 | attached to the function's IAM role: 105 | 106 | ```json 107 | { 108 | "Version": "2012-10-17", 109 | "Statement": [ 110 | { 111 | "Effect": "Allow", 112 | "Action": [ 113 | "lambda:AddPermission", 114 | "lambda:RemovePermission", 115 | "events:PutRule", 116 | "events:DeleteRule", 117 | "events:PutTargets", 118 | "events:RemoveTargets" 119 | ], 120 | "Resource": "*" 121 | } 122 | ] 123 | } 124 | ``` 125 | ### Certificate Verification 126 | To turn off certification verification, or to use a custom CA bundle path for the underlying boto3 clients used by this library, override the `ssl_verify` argument with the appropriate values. These can be either: 127 | * `False` - do not validate SSL certificates. SSL will still be used, but SSL certificates will not be verified. 128 | * `path/to/cert/bundle.pem` - A filename of the CA cert bundle to uses. You can specify this argument if you want to use a different CA cert bundle than the one used by botocore. 129 | 130 | ### Use CDK to depoy a Custom Resource that uses Custom Resource Helper 131 | 132 | You can use the [AWS Cloud Development Kit (AWS CDK)](https://docs.aws.amazon.com/cdk/v2/guide/home.html) to deploy a Custom Resource that uses Custom Resource Helper. AWS CDK is an open-source software development framework for defining cloud infrastructure in code and provisioning it through AWS CloudFormation. 133 | 134 | **Note**: `crhelper` is not intended to be used with AWS CDK using the [Provider](https://docs.aws.amazon.com/cdk/api/v2/docs/aws-cdk-lib.custom_resources.Provider.html) construct. 135 | 136 | #### AWS CDK template example 137 | ``` 138 | from aws_cdk import ( 139 | ... 140 | aws_lambda as _lambda, 141 | CustomResource, 142 | ) 143 | 144 | crhelperSumResource = _lambda.Function(...) 145 | 146 | customResource = CustomResource( 147 | self, 148 | 'MyCustomResource' 149 | serviceToken = crhelperSumResource.function_arn, 150 | properties = { 151 | 'No1': 1, 152 | 'No2': 2 153 | }, 154 | ) 155 | 156 | 157 | ``` 158 | 159 | ## Credits 160 | 161 | Decorator implementation inspired by https://github.com/ryansb/cfn-wrapper-python 162 | 163 | Log implementation inspired by https://gitlab.com/hadrien/aws_lambda_logging 164 | 165 | ## License 166 | 167 | This library is licensed under the Apache 2.0 License. 168 | -------------------------------------------------------------------------------- /crhelper/__init__.py: -------------------------------------------------------------------------------- 1 | from crhelper.resource_helper import CfnResource, SUCCESS, FAILED 2 | -------------------------------------------------------------------------------- /crhelper/__init__.pyi: -------------------------------------------------------------------------------- 1 | """ 2 | *.pyi files are auto-generated with the following commands: 3 | 4 | $> pip install mypy 5 | $> stubgen ./crhelper -o . 6 | """ 7 | from crhelper.resource_helper import CfnResource as CfnResource, FAILED as FAILED, SUCCESS as SUCCESS 8 | -------------------------------------------------------------------------------- /crhelper/log_helper.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import json 3 | import logging 4 | 5 | 6 | def _json_formatter(obj): 7 | """Formatter for unserialisable values.""" 8 | return str(obj) 9 | 10 | 11 | class JsonFormatter(logging.Formatter): 12 | """AWS Lambda Logging formatter. 13 | 14 | Formats the log message as a JSON encoded string. If the message is a 15 | dict it will be used directly. If the message can be parsed as JSON, then 16 | the parse d value is used in the output record. 17 | """ 18 | 19 | def __init__(self, **kwargs): 20 | super(JsonFormatter, self).__init__() 21 | self.format_dict = { 22 | 'timestamp': '%(asctime)s', 23 | 'level': '%(levelname)s', 24 | 'location': '%(name)s.%(funcName)s:%(lineno)d', 25 | } 26 | self.format_dict.update(kwargs) 27 | self.default_json_formatter = kwargs.pop( 28 | 'json_default', _json_formatter) 29 | 30 | def format(self, record): 31 | record_dict = record.__dict__.copy() 32 | record_dict['asctime'] = self.formatTime(record) 33 | 34 | log_dict = { 35 | k: v % record_dict 36 | for k, v in self.format_dict.items() 37 | if v 38 | } 39 | 40 | if isinstance(record_dict['msg'], dict): 41 | log_dict['message'] = record_dict['msg'] 42 | else: 43 | log_dict['message'] = record.getMessage() 44 | 45 | # Attempt to decode the message as JSON, if so, merge it with the 46 | # overall message for clarity. 47 | try: 48 | log_dict['message'] = json.loads(log_dict['message']) 49 | except (TypeError, ValueError): 50 | pass 51 | 52 | if record.exc_info: 53 | # Cache the traceback text to avoid converting it multiple times 54 | # (it's constant anyway) 55 | # from logging.Formatter:format 56 | if not record.exc_text: 57 | record.exc_text = self.formatException(record.exc_info) 58 | 59 | if record.exc_text: 60 | log_dict['exception'] = record.exc_text 61 | 62 | json_record = json.dumps(log_dict, default=self.default_json_formatter) 63 | 64 | if hasattr(json_record, 'decode'): # pragma: no cover 65 | json_record = json_record.decode('utf-8') 66 | 67 | return json_record 68 | 69 | 70 | def setupLogger(level='DEBUG', formatter_cls=JsonFormatter, boto_level=None, **kwargs): 71 | if formatter_cls: 72 | for handler in logging.root.handlers: 73 | handler.setFormatter(formatter_cls(**kwargs)) 74 | 75 | logging.root.setLevel(level=level) 76 | 77 | if not boto_level: 78 | boto_level = level 79 | 80 | logging.getLogger('boto').setLevel(boto_level) 81 | logging.getLogger('boto3').setLevel(boto_level) 82 | logging.getLogger('botocore').setLevel(boto_level) 83 | logging.getLogger('urllib3').setLevel(boto_level) 84 | -------------------------------------------------------------------------------- /crhelper/log_helper.pyi: -------------------------------------------------------------------------------- 1 | import logging 2 | from typing import Any, Optional 3 | 4 | class JsonFormatter(logging.Formatter): 5 | format_dict: Any = ... 6 | default_json_formatter: Any = ... 7 | def __init__(self, **kwargs: Any) -> None: ... 8 | def format(self, record: Any): ... 9 | 10 | def setupLogger(level: str = ..., formatter_cls: Any = ..., boto_level: Optional[Any] = ..., **kwargs: Any) -> None: ... 11 | -------------------------------------------------------------------------------- /crhelper/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-cloudformation/custom-resource-helper/b07bd4ec17e175a5148ee07e7d5efa12dd7bb6a5/crhelper/py.typed -------------------------------------------------------------------------------- /crhelper/resource_helper.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | """ 3 | TODO: 4 | * Async mode – take a wait condition handle as an input, increases max timeout to 12 hours 5 | * Idempotency – If a duplicate request comes in (say there was a network error in signaling back to cfn) the subsequent 6 | request should return the already created response, will need a persistent store of some kind... 7 | * Functional tests 8 | """ 9 | 10 | from __future__ import print_function 11 | import threading 12 | from crhelper.utils import _send_response 13 | from crhelper import log_helper 14 | import logging 15 | import random 16 | import boto3 17 | import string 18 | import json 19 | import os 20 | from time import sleep 21 | 22 | logger = logging.getLogger(__name__) 23 | 24 | SUCCESS = 'SUCCESS' 25 | FAILED = 'FAILED' 26 | 27 | 28 | class CfnResource(object): 29 | 30 | def __init__(self, json_logging=False, log_level='DEBUG', boto_level='ERROR', polling_interval=2, sleep_on_delete=120, ssl_verify=None): 31 | self._sleep_on_delete = sleep_on_delete 32 | self._create_func = None 33 | self._update_func = None 34 | self._delete_func = None 35 | self._poll_create_func = None 36 | self._poll_update_func = None 37 | self._poll_delete_func = None 38 | self._timer = None 39 | self._init_failed = None 40 | self._json_logging = json_logging 41 | self._log_level = log_level 42 | self._boto_level = boto_level 43 | self._send_response = False 44 | self._polling_interval = polling_interval 45 | self.Status = "" 46 | self.Reason = "" 47 | self.PhysicalResourceId = "" 48 | self.StackId = "" 49 | self.RequestId = "" 50 | self.LogicalResourceId = "" 51 | self.Data = {} 52 | self.NoEcho = False 53 | self._event = {} 54 | self._context = None 55 | self._response_url = "" 56 | self._sam_local = os.getenv('AWS_SAM_LOCAL') 57 | self._region = os.getenv('AWS_REGION') 58 | self._ssl_verify = ssl_verify 59 | try: 60 | if not self._sam_local: 61 | self._lambda_client = boto3.client('lambda', region_name=self._region, verify=self._ssl_verify) 62 | self._events_client = boto3.client('events', region_name=self._region, verify=self._ssl_verify) 63 | self._logs_client = boto3.client('logs', region_name=self._region, verify=self._ssl_verify) 64 | if json_logging: 65 | log_helper.setupLogger(log_level, boto_level=boto_level, RequestType='ContainerInit') 66 | else: 67 | log_helper.setupLogger(log_level, formatter_cls=None, boto_level=boto_level) 68 | except Exception as e: 69 | logger.error(e, exc_info=True) 70 | self.init_failure(e) 71 | 72 | def __call__(self, event, context): 73 | try: 74 | self._log_setup(event, context) 75 | logger.debug(event) 76 | if not self._crhelper_init(event, context): 77 | return 78 | # Check for polling functions 79 | if self._poll_enabled() and self._sam_local: 80 | logger.info("Skipping poller functionality, as this is a local invocation") 81 | elif self._poll_enabled(): 82 | self._polling_init(event) 83 | # If polling is not enabled, then we should respond 84 | else: 85 | logger.debug("enabling send_response") 86 | self._send_response = True 87 | logger.debug("_send_response: %s" % self._send_response) 88 | if self._send_response: 89 | if self.RequestType == 'Delete': 90 | self._wait_for_cwlogs() 91 | self._cfn_response(event) 92 | except Exception as e: 93 | logger.error(e, exc_info=True) 94 | self._send(FAILED, str(e)) 95 | finally: 96 | if self._timer: 97 | self._timer.cancel() 98 | 99 | def _wait_for_cwlogs(self, sleep=sleep): 100 | time_left = int(self._context.get_remaining_time_in_millis() / 1000) - 15 101 | sleep_time = 0 102 | 103 | if time_left > self._sleep_on_delete: 104 | sleep_time = self._sleep_on_delete 105 | 106 | if sleep_time > 1: 107 | sleep(sleep_time) 108 | 109 | def _log_setup(self, event, context): 110 | if self._json_logging: 111 | log_helper.setupLogger(self._log_level, boto_level=self._boto_level, RequestType=event['RequestType'], 112 | StackId=event['StackId'], RequestId=event['RequestId'], 113 | LogicalResourceId=event['LogicalResourceId'], aws_request_id=context.aws_request_id) 114 | else: 115 | log_helper.setupLogger(self._log_level, boto_level=self._boto_level, formatter_cls=None) 116 | 117 | def _crhelper_init(self, event, context): 118 | self._send_response = False 119 | self.Status = SUCCESS 120 | self.Reason = "" 121 | self.PhysicalResourceId = "" 122 | self.StackId = event["StackId"] 123 | self.RequestId = event["RequestId"] 124 | self.LogicalResourceId = event["LogicalResourceId"] 125 | self.Data = {} 126 | if "CrHelperData" in event.keys(): 127 | self.Data = event["CrHelperData"] 128 | self.RequestType = event["RequestType"] 129 | self._event = event 130 | self._context = context 131 | self._response_url = event['ResponseURL'] 132 | if self._timer: 133 | self._timer.cancel() 134 | if self._init_failed: 135 | self._send(FAILED, str(self._init_failed)) 136 | return False 137 | self._set_timeout() 138 | self._wrap_function(self._get_func()) 139 | return True 140 | 141 | def _polling_init(self, event): 142 | # Setup polling on initial request 143 | logger.debug("pid1: %s" % self.PhysicalResourceId) 144 | if 'CrHelperPoll' not in event.keys() and self.Status != FAILED: 145 | logger.info("Setting up polling") 146 | self.Data["PhysicalResourceId"] = self.PhysicalResourceId 147 | self._setup_polling() 148 | self.PhysicalResourceId = None 149 | logger.debug("pid2: %s" % self.PhysicalResourceId) 150 | # if physical id is set, or there was a failure then we're done 151 | logger.debug("pid3: %s" % self.PhysicalResourceId) 152 | if self.PhysicalResourceId or self.Status == FAILED: 153 | logger.info("Polling complete, removing cwe schedule") 154 | self._remove_polling() 155 | self._send_response = True 156 | 157 | def generate_physical_id(self, event): 158 | return '_'.join([ 159 | event['StackId'].split('/')[1], 160 | event['LogicalResourceId'], 161 | self._rand_string(8) 162 | ]) 163 | 164 | def _cfn_response(self, event): 165 | # Use existing PhysicalResourceId if it's in the event and no ID was set 166 | if not self.PhysicalResourceId and "PhysicalResourceId" in event.keys(): 167 | logger.info("PhysicalResourceId present in event, Using that for response") 168 | self.PhysicalResourceId = event['PhysicalResourceId'] 169 | # Generate a physical id if none is provided 170 | elif not self.PhysicalResourceId or self.PhysicalResourceId is True: 171 | logger.info("No physical resource id returned, generating one...") 172 | self.PhysicalResourceId = self.generate_physical_id(event) 173 | self._send() 174 | 175 | def _poll_enabled(self): 176 | return getattr(self, "_poll_{}_func".format(self._event['RequestType'].lower())) 177 | 178 | def create(self, func): 179 | self._create_func = func 180 | return func 181 | 182 | def update(self, func): 183 | self._update_func = func 184 | return func 185 | 186 | def delete(self, func): 187 | self._delete_func = func 188 | return func 189 | 190 | def poll_create(self, func): 191 | self._poll_create_func = func 192 | return func 193 | 194 | def poll_update(self, func): 195 | self._poll_update_func = func 196 | return func 197 | 198 | def poll_delete(self, func): 199 | self._poll_delete_func = func 200 | return func 201 | 202 | def _wrap_function(self, func): 203 | try: 204 | self.PhysicalResourceId = func(self._event, self._context) if func else '' 205 | except Exception as e: 206 | logger.error(str(e), exc_info=True) 207 | self.Reason = str(e) 208 | self.Status = FAILED 209 | 210 | def _timeout(self): 211 | logger.error("Execution is about to time out, sending failure message") 212 | self._send(FAILED, "Execution timed out") 213 | 214 | def _set_timeout(self): 215 | self._timer = threading.Timer((self._context.get_remaining_time_in_millis() / 1000.00) - 0.5, 216 | self._timeout) 217 | self._timer.start() 218 | 219 | def _get_func(self): 220 | request_type = "_{}_func" 221 | if "CrHelperPoll" in self._event.keys(): 222 | request_type = "_poll" + request_type 223 | return getattr(self, request_type.format(self._event['RequestType'].lower())) 224 | 225 | def _send(self, status=None, reason="", send_response=_send_response): 226 | if len(str(str(self.Reason))) > 256: 227 | self.Reason = "ERROR: (truncated) " + str(self.Reason)[len(str(self.Reason)) - 240:] 228 | if len(str(reason)) > 256: 229 | reason = "ERROR: (truncated) " + str(reason)[len(str(reason)) - 240:] 230 | response_body = { 231 | 'Status': self.Status, 232 | 'PhysicalResourceId': str(self.PhysicalResourceId), 233 | 'StackId': self.StackId, 234 | 'RequestId': self.RequestId, 235 | 'LogicalResourceId': self.LogicalResourceId, 236 | 'Reason': str(self.Reason), 237 | 'Data': self.Data, 238 | 'NoEcho': self.NoEcho, 239 | } 240 | if status: 241 | response_body.update({'Status': status, 'Reason': reason}) 242 | send_response(self._response_url, response_body, self._ssl_verify) 243 | 244 | def init_failure(self, error): 245 | self._init_failed = error 246 | logger.error(str(error), exc_info=True) 247 | 248 | def _cleanup_response(self): 249 | for k in ["CrHelperPoll", "CrHelperPermission", "CrHelperRule"]: 250 | if k in self.Data.keys(): 251 | del self.Data[k] 252 | 253 | @staticmethod 254 | def _rand_string(l): 255 | return ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(l)) 256 | 257 | def _add_permission(self, rule_arn): 258 | sid = self._event['LogicalResourceId'] + self._rand_string(8) 259 | self._lambda_client.add_permission( 260 | FunctionName=self._context.function_name, 261 | StatementId=sid, 262 | Action='lambda:InvokeFunction', 263 | Principal='events.amazonaws.com', 264 | SourceArn=rule_arn 265 | ) 266 | return sid 267 | 268 | def _put_rule(self): 269 | schedule_unit = 'minutes' if self._polling_interval != 1 else 'minute' 270 | response = self._events_client.put_rule( 271 | Name=self._event['LogicalResourceId'] + self._rand_string(8), 272 | ScheduleExpression='rate({} {})'.format(self._polling_interval, schedule_unit), 273 | State='ENABLED', 274 | ) 275 | return response["RuleArn"] 276 | 277 | def _put_targets(self, func_name): 278 | region = self._event['CrHelperRule'].split(":")[3] 279 | account_id = self._event['CrHelperRule'].split(":")[4] 280 | partition = self._event['CrHelperRule'].split(":")[1] 281 | rule_name = self._event['CrHelperRule'].split("/")[1] 282 | logger.debug(self._event) 283 | self._events_client.put_targets( 284 | Rule=rule_name, 285 | Targets=[ 286 | { 287 | 'Id': '1', 288 | 'Arn': 'arn:%s:lambda:%s:%s:function:%s' % (partition, region, account_id, func_name), 289 | 'Input': json.dumps(self._event) 290 | } 291 | ] 292 | ) 293 | 294 | def _remove_targets(self, rule_arn): 295 | self._events_client.remove_targets( 296 | Rule=rule_arn.split("/")[1], 297 | Ids=['1'] 298 | ) 299 | 300 | def _remove_permission(self, sid): 301 | self._lambda_client.remove_permission( 302 | FunctionName=self._context.function_name, 303 | StatementId=sid 304 | ) 305 | 306 | def _delete_rule(self, rule_arn): 307 | self._events_client.delete_rule( 308 | Name=rule_arn.split("/")[1] 309 | ) 310 | 311 | def _setup_polling(self): 312 | self._event['CrHelperData'] = self.Data 313 | self._event['CrHelperPoll'] = True 314 | self._event['CrHelperRule'] = self._put_rule() 315 | self._event['CrHelperPermission'] = self._add_permission(self._event['CrHelperRule']) 316 | self._put_targets(self._context.function_name) 317 | 318 | def _remove_polling(self): 319 | if 'CrHelperData' in self._event.keys(): 320 | self._event.pop('CrHelperData') 321 | if "PhysicalResourceId" in self.Data.keys(): 322 | self.Data.pop("PhysicalResourceId") 323 | if 'CrHelperRule' in self._event.keys(): 324 | self._remove_targets(self._event['CrHelperRule']) 325 | else: 326 | logger.error("Cannot remove CloudWatch events rule, Rule arn not available in event") 327 | if 'CrHelperPermission' in self._event.keys(): 328 | self._remove_permission(self._event['CrHelperPermission']) 329 | else: 330 | logger.error("Cannot remove lambda events permission, permission id not available in event") 331 | if 'CrHelperRule' in self._event.keys(): 332 | self._delete_rule(self._event['CrHelperRule']) 333 | else: 334 | logger.error("Cannot remove CloudWatch events target, Rule arn not available in event") 335 | -------------------------------------------------------------------------------- /crhelper/resource_helper.pyi: -------------------------------------------------------------------------------- 1 | from crhelper import log_helper as log_helper 2 | from typing import Any, Callable, Optional, Union 3 | 4 | logger: Any 5 | SUCCESS: str 6 | FAILED: str 7 | 8 | _DecoratorTypeDef = Callable[[Any, Any], Union[bool, str, None]] 9 | 10 | class CfnResource: 11 | Status: str = ... 12 | Reason: str = ... 13 | PhysicalResourceId: str = ... 14 | StackId: str = ... 15 | RequestId: str = ... 16 | LogicalResourceId: str = ... 17 | Data: Any = ... 18 | NoEcho: bool = ... 19 | def __init__(self, json_logging: bool = ..., log_level: str = ..., boto_level: str = ..., polling_interval: int = ..., sleep_on_delete: int = ..., ssl_verify: Optional[bool] = ...) -> None: ... 20 | def __call__(self, event: Any, context: Any) -> None: ... 21 | def generate_physical_id(self, event: Any): ... 22 | def create(self, func: _DecoratorTypeDef) -> _DecoratorTypeDef: ... 23 | def update(self, func: _DecoratorTypeDef) -> _DecoratorTypeDef: ... 24 | def delete(self, func: _DecoratorTypeDef) -> _DecoratorTypeDef: ... 25 | def poll_create(self, func: _DecoratorTypeDef) -> _DecoratorTypeDef: ... 26 | def poll_update(self, func: _DecoratorTypeDef) -> _DecoratorTypeDef: ... 27 | def poll_delete(self, func: _DecoratorTypeDef) -> _DecoratorTypeDef: ... 28 | def init_failure(self, error: Any) -> None: ... 29 | -------------------------------------------------------------------------------- /crhelper/utils.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | 3 | import json 4 | import logging as logging 5 | import ssl 6 | import time 7 | from http.client import HTTPSConnection 8 | from os import path 9 | from typing import Union, AnyStr 10 | from urllib.parse import urlsplit, urlunsplit 11 | 12 | logger = logging.getLogger(__name__) 13 | MAX_RETRIES = 5 # Maximum number of retries 14 | 15 | def _send_response(response_url: AnyStr, response_body: AnyStr, ssl_verify: Union[bool, AnyStr] = None): 16 | try: 17 | json_response_body = json.dumps(response_body) 18 | except Exception as e: 19 | msg = "Failed to convert response to json: {}".format(str(e)) 20 | logger.error(msg, exc_info=True) 21 | response_body = {'Status': 'FAILED', 'Data': {}, 'Reason': msg} 22 | json_response_body = json.dumps(response_body) 23 | logger.debug("CFN response URL: {}".format(response_url)) 24 | logger.debug(json_response_body) 25 | headers = {'content-type': '', 'content-length': str(len(json_response_body))} 26 | split_url = urlsplit(response_url) 27 | host = split_url.netloc 28 | url = urlunsplit(("", "", *split_url[2:])) 29 | ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) 30 | if isinstance(ssl_verify, str): 31 | if path.exists(ssl_verify): 32 | ctx.load_verify_locations(cafile=ssl_verify) 33 | else: 34 | logger.warning("Cert path {0} does not exist!. Falling back to using system cafile.".format(ssl_verify)) 35 | if ssl_verify is False: 36 | ctx.check_hostname = False 37 | ctx.verify_mode = ssl.CERT_NONE 38 | # If ssl_verify is True or None dont modify the context in any way. 39 | 40 | while True: 41 | retry_count = 0 42 | success = False 43 | while retry_count < MAX_RETRIES and not success: 44 | try: 45 | connection = HTTPSConnection(host, context=ctx) 46 | connection.request(method="PUT", url=url, body=json_response_body, headers=headers) 47 | response = connection.getresponse() 48 | logger.info("CloudFormation returned status code: {}".format(response.reason)) 49 | success = True 50 | except Exception as e: 51 | retry_count += 1 52 | logger.error("Unexpected failure sending response to CloudFormation {}. Retrying in 2 seconds...".format(e), exc_info=True) 53 | time.sleep(2) 54 | if success: 55 | break 56 | else: 57 | logger.error("Maximum retries reached. Unable to send response to CloudFormation.") 58 | time.sleep(5) -------------------------------------------------------------------------------- /crhelper/utils.pyi: -------------------------------------------------------------------------------- 1 | from typing import Any 2 | 3 | logger: Any 4 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | boto3>=1.9.108 2 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | 2 | """ 3 | Copyright 2016-2018 Amazon.com, Inc. or its affiliates. All Rights Reserved. 4 | Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance with the 5 | License. A copy of the License is located at 6 | http://aws.amazon.com/apache2.0/ 7 | or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR 8 | CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and 9 | limitations under the License. 10 | """ 11 | 12 | from setuptools import setup, find_packages 13 | 14 | 15 | with open("README.md", "r") as fh: 16 | long_description = fh.read() 17 | 18 | 19 | setup( 20 | name="crhelper", 21 | version="2.0.12", 22 | description="crhelper simplifies authoring CloudFormation Custom Resources", 23 | long_description=long_description, 24 | long_description_content_type="text/markdown", 25 | url="https://github.com/aws-cloudformation/custom-resource-helper", 26 | author="Jay McConnell", 27 | author_email="jmmccon@amazon.com", 28 | license="Apache2", 29 | packages=find_packages(exclude=("tests", "tests.*",)), 30 | package_data={"crhelper": ["py.typed", "*.pyi"]}, 31 | install_requires=[], 32 | tests_require=["boto3"], 33 | test_suite="tests", 34 | classifiers=[ 35 | 'Programming Language :: Python :: 3.6', 36 | 'Programming Language :: Python :: 3.7', 37 | 'Programming Language :: Python :: 3.8', 38 | 'Programming Language :: Python :: 3.9', 39 | 'Programming Language :: Python :: 3.10', 40 | "License :: OSI Approved :: Apache Software License", 41 | "Operating System :: OS Independent", 42 | ], 43 | ) 44 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-cloudformation/custom-resource-helper/b07bd4ec17e175a5148ee07e7d5efa12dd7bb6a5/tests/__init__.py -------------------------------------------------------------------------------- /tests/functional/create.json: -------------------------------------------------------------------------------- 1 | { 2 | "RequestType": "Create", 3 | "ResponseURL": "https://s3.amazonaws.com/test-bucket/test-object", 4 | "StackId": "arn:{partition}:cloudformation:{region}:EXAMPLE/stack-name/guid", 5 | "RequestId": "unique id for this create request", 6 | "ResourceType": "Custom::TestResource", 7 | "LogicalResourceId": "MyTestResource", 8 | "ResourceProperties": { 9 | "StackName": "stack-name", 10 | "List": [ 11 | "1", 12 | "2", 13 | "3" 14 | ] 15 | } 16 | } -------------------------------------------------------------------------------- /tests/functional/lambda_function.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from crhelper import CfnResource 3 | import logging 4 | import boto3 5 | 6 | logger = logging.getLogger(__name__) 7 | # Initialise the helper, all inputs are optional, this example shows the defaults 8 | helper = CfnResource(json_logging=False, log_level='DEBUG', boto_level='CRITICAL') 9 | 10 | try: 11 | ## Init code goes here 12 | # secrets_client = boto3.client('secretsmanager') 13 | pass 14 | except Exception as e: 15 | helper.init_failure(e) 16 | 17 | 18 | @helper.create 19 | def create(event, context): 20 | logger.info("Got Create") 21 | # Optionally return an ID that will be used for the resource PhysicalResourceId, if None is returned an ID will be 22 | # generated. If a poll_create function is defined return value is ignored 23 | return "MyResourceId" 24 | 25 | 26 | @helper.update 27 | def update(event, context): 28 | logger.info("Got Update") 29 | # If the update resulted in a new resource being created, return an id for the new resource. CloudFormation will send 30 | # a delete event with the old id when stack update completes 31 | 32 | 33 | @helper.delete 34 | def delete(event, context): 35 | logger.info("Got Delete") 36 | # Delete never returns anything. Should not fail if the underlying resources are already deleted. Desired state. 37 | 38 | 39 | @helper.poll_create 40 | def poll_create(event, context): 41 | logger.info("Got create poll") 42 | # Return a resource id or True to indicate that creation is complete. if True is returned an id will be generated 43 | return True 44 | 45 | 46 | def handler(event, context): 47 | helper(event, context) 48 | -------------------------------------------------------------------------------- /tests/test_log_helper.py: -------------------------------------------------------------------------------- 1 | from crhelper.log_helper import * 2 | import unittest 3 | import logging 4 | 5 | 6 | class TestLogHelper(unittest.TestCase): 7 | 8 | def test_logging_no_formatting(self): 9 | logger = logging.getLogger('1') 10 | handler = logging.StreamHandler() 11 | logger.addHandler(handler) 12 | orig_formatters = [] 13 | for c in range(len(logging.root.handlers)): 14 | orig_formatters.append(logging.root.handlers[c].formatter) 15 | setupLogger(level='DEBUG', formatter_cls=None, boto_level='CRITICAL') 16 | new_formatters = [] 17 | for c in range(len(logging.root.handlers)): 18 | new_formatters.append(logging.root.handlers[c].formatter) 19 | self.assertEqual(orig_formatters, new_formatters) 20 | 21 | def test_logging_boto_explicit(self): 22 | logger = logging.getLogger('2') 23 | handler = logging.StreamHandler() 24 | logger.addHandler(handler) 25 | setupLogger(level='DEBUG', formatter_cls=None, boto_level='CRITICAL') 26 | for t in ['boto', 'boto3', 'botocore', 'urllib3']: 27 | b_logger = logging.getLogger(t) 28 | self.assertEqual(b_logger.level, 50) 29 | 30 | def test_logging_json(self): 31 | logger = logging.getLogger('3') 32 | handler = logging.StreamHandler() 33 | logger.addHandler(handler) 34 | setupLogger(level='DEBUG', formatter_cls=JsonFormatter, RequestType='ContainerInit') 35 | for handler in logging.root.handlers: 36 | self.assertEqual(JsonFormatter, type(handler.formatter)) 37 | 38 | def test_logging_boto_implicit(self): 39 | logger = logging.getLogger('4') 40 | handler = logging.StreamHandler() 41 | logger.addHandler(handler) 42 | setupLogger(level='DEBUG', formatter_cls=JsonFormatter, RequestType='ContainerInit') 43 | for t in ['boto', 'boto3', 'botocore', 'urllib3']: 44 | b_logger = logging.getLogger(t) 45 | self.assertEqual(b_logger.level, 10) 46 | 47 | def test_logging_json_keys(self): 48 | with self.assertLogs() as ctx: 49 | logger = logging.getLogger() 50 | handler = logging.StreamHandler() 51 | logger.addHandler(handler) 52 | setupLogger(level='DEBUG', formatter_cls=JsonFormatter, RequestType='ContainerInit') 53 | logger.info("test") 54 | logs = json.loads(ctx.output[0]) 55 | self.assertEqual(["timestamp", "level", "location", "RequestType", "message"], list(logs.keys())) 56 | 57 | def test_logging_json_parse_message(self): 58 | with self.assertLogs() as ctx: 59 | logger = logging.getLogger() 60 | handler = logging.StreamHandler() 61 | logger.addHandler(handler) 62 | setupLogger(level='DEBUG', formatter_cls=JsonFormatter, RequestType='ContainerInit') 63 | logger.info("{}") 64 | logs = json.loads(ctx.output[0]) 65 | self.assertEqual({}, logs["message"]) 66 | 67 | def test_logging_json_exception(self): 68 | with self.assertLogs() as ctx: 69 | logger = logging.getLogger() 70 | handler = logging.StreamHandler() 71 | logger.addHandler(handler) 72 | setupLogger(level='DEBUG', formatter_cls=JsonFormatter, RequestType='ContainerInit') 73 | try: 74 | 1 + 't' 75 | except Exception as e: 76 | logger.info("[]", exc_info=True) 77 | logs = json.loads(ctx.output[0]) 78 | self.assertIn("exception", logs.keys()) 79 | -------------------------------------------------------------------------------- /tests/test_resource_helper.py: -------------------------------------------------------------------------------- 1 | import os 2 | import crhelper 3 | import unittest 4 | from unittest.mock import call, patch, Mock 5 | import threading 6 | 7 | test_events = { 8 | "Create": { 9 | "RequestType": "Create", 10 | "RequestId": "test-event-id", 11 | "StackId": "arn/test-stack-id/guid", 12 | "LogicalResourceId": "TestResourceId", 13 | "ResponseURL": "response_url" 14 | }, 15 | "Update": { 16 | "RequestType": "Update", 17 | "RequestId": "test-event-id", 18 | "StackId": "test-stack-id", 19 | "LogicalResourceId": "TestResourceId", 20 | "PhysicalResourceId": "test-pid", 21 | "ResponseURL": "response_url" 22 | }, 23 | "Delete": { 24 | "RequestType": "Delete", 25 | "RequestId": "test-event-id", 26 | "StackId": "test-stack-id", 27 | "LogicalResourceId": "TestResourceId", 28 | "PhysicalResourceId": "test-pid", 29 | "ResponseURL": "response_url" 30 | } 31 | } 32 | 33 | 34 | class MockContext(object): 35 | 36 | function_name = "test-function" 37 | ms_remaining = 9000 38 | 39 | @staticmethod 40 | def get_remaining_time_in_millis(): 41 | return MockContext.ms_remaining 42 | 43 | 44 | class TestCfnResource(unittest.TestCase): 45 | def setUp(self): 46 | os.environ['AWS_REGION'] = 'us-east-1' 47 | 48 | def tearDown(self): 49 | os.environ.pop('AWS_REGION', None) 50 | 51 | @patch('crhelper.log_helper.setupLogger', return_value=None) 52 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 53 | def test_init(self, mock_method): 54 | crhelper.resource_helper.CfnResource() 55 | mock_method.assert_called_once_with('DEBUG', boto_level='ERROR', formatter_cls=None) 56 | 57 | crhelper.resource_helper.CfnResource(json_logging=True) 58 | mock_method.assert_called_with('DEBUG', boto_level='ERROR', RequestType='ContainerInit') 59 | 60 | @patch('crhelper.log_helper.setupLogger', return_value=None) 61 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 62 | def test_init_failure(self, mock_method): 63 | mock_method.side_effect = Exception("test") 64 | c = crhelper.resource_helper.CfnResource(json_logging=True) 65 | self.assertTrue(c._init_failed) 66 | 67 | @patch('crhelper.log_helper.setupLogger', Mock()) 68 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 69 | @patch('crhelper.resource_helper.CfnResource._polling_init', Mock()) 70 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 71 | @patch('crhelper.resource_helper.CfnResource._send') 72 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 73 | @patch('crhelper.resource_helper.CfnResource._wrap_function', Mock()) 74 | def test_init_failure_call(self, mock_send): 75 | c = crhelper.resource_helper.CfnResource() 76 | c.init_failure(Exception('TestException')) 77 | 78 | event = test_events["Create"] 79 | c.__call__(event, MockContext) 80 | 81 | self.assertEqual([call('FAILED', 'TestException')], mock_send.call_args_list) 82 | 83 | @patch('crhelper.log_helper.setupLogger', Mock()) 84 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 85 | @patch('crhelper.resource_helper.CfnResource._polling_init', Mock()) 86 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 87 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 88 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 89 | @patch('crhelper.resource_helper.CfnResource._wrap_function', Mock()) 90 | @patch('crhelper.resource_helper.CfnResource._cfn_response', return_value=None) 91 | def test_call(self, cfn_response_mock): 92 | c = crhelper.resource_helper.CfnResource() 93 | event = test_events["Create"] 94 | c.__call__(event, MockContext) 95 | self.assertTrue(c._send_response) 96 | cfn_response_mock.assert_called_once_with(event) 97 | 98 | c._sam_local = True 99 | c._poll_enabled = Mock(return_value=True) 100 | c._polling_init = Mock() 101 | c.__call__(event, MockContext) 102 | c._polling_init.assert_not_called() 103 | self.assertEqual(1, len(cfn_response_mock.call_args_list)) 104 | 105 | c._sam_local = False 106 | c._send_response = False 107 | c.__call__(event, MockContext) 108 | c._polling_init.assert_called() 109 | self.assertEqual(1, len(cfn_response_mock.call_args_list)) 110 | 111 | event = test_events["Delete"] 112 | c._wait_for_cwlogs = Mock() 113 | c._poll_enabled = Mock(return_value=False) 114 | c.__call__(event, MockContext) 115 | c._wait_for_cwlogs.assert_called() 116 | 117 | c._send = Mock() 118 | cfn_response_mock.side_effect = Exception("test") 119 | c.__call__(event, MockContext) 120 | c._send.assert_called_with('FAILED', "test") 121 | 122 | @patch('crhelper.log_helper.setupLogger', Mock()) 123 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 124 | @patch('crhelper.resource_helper.CfnResource._polling_init', Mock()) 125 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 126 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 127 | @patch('crhelper.resource_helper.CfnResource._wrap_function', Mock()) 128 | @patch('crhelper.resource_helper.CfnResource._cfn_response', Mock(return_value=None)) 129 | def test_wait_for_cwlogs(self): 130 | 131 | c = crhelper.resource_helper.CfnResource() 132 | c._context = MockContext 133 | s = Mock() 134 | c._wait_for_cwlogs(sleep=s) 135 | s.assert_not_called() 136 | MockContext.ms_remaining = 140000 137 | c._wait_for_cwlogs(sleep=s) 138 | s.assert_called_once() 139 | 140 | @patch('crhelper.log_helper.setupLogger', Mock()) 141 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 142 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 143 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 144 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 145 | @patch('crhelper.resource_helper.CfnResource._wrap_function', Mock()) 146 | @patch('crhelper.resource_helper.CfnResource._cfn_response', Mock()) 147 | def test_polling_init(self): 148 | c = crhelper.resource_helper.CfnResource() 149 | event = test_events['Create'] 150 | c._setup_polling = Mock() 151 | c._remove_polling = Mock() 152 | c._polling_init(event) 153 | c._setup_polling.assert_called_once() 154 | c._remove_polling.assert_not_called() 155 | self.assertEqual(c.PhysicalResourceId, None) 156 | 157 | c.Status = 'FAILED' 158 | c._setup_polling.assert_called_once() 159 | c._setup_polling.assert_called_once() 160 | 161 | c = crhelper.resource_helper.CfnResource() 162 | event = test_events['Create'] 163 | c._setup_polling = Mock() 164 | c._remove_polling = Mock() 165 | event['CrHelperPoll'] = "Some stuff" 166 | c.PhysicalResourceId = None 167 | c._polling_init(event) 168 | c._remove_polling.assert_not_called() 169 | c._setup_polling.assert_not_called() 170 | 171 | c.Status = 'FAILED' 172 | c._polling_init(event) 173 | c._remove_polling.assert_called_once() 174 | c._setup_polling.assert_not_called() 175 | 176 | c.Status = '' 177 | c.PhysicalResourceId = "some-id" 178 | c._remove_polling.assert_called() 179 | c._setup_polling.assert_not_called() 180 | 181 | @patch('crhelper.log_helper.setupLogger', Mock()) 182 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 183 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 184 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 185 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 186 | @patch('crhelper.resource_helper.CfnResource._wrap_function', Mock()) 187 | def test_cfn_response(self): 188 | c = crhelper.resource_helper.CfnResource() 189 | event = test_events['Create'] 190 | c._send = Mock() 191 | 192 | orig_pid = c.PhysicalResourceId 193 | self.assertEqual(orig_pid, '') 194 | c._cfn_response(event) 195 | c._send.assert_called_once() 196 | print("RID: [%s]" % [c.PhysicalResourceId]) 197 | self.assertEqual(True, c.PhysicalResourceId.startswith('test-stack-id_TestResourceId_')) 198 | 199 | c._send = Mock() 200 | c.PhysicalResourceId = 'testpid' 201 | c._cfn_response(event) 202 | c._send.assert_called_once() 203 | self.assertEqual('testpid', c.PhysicalResourceId) 204 | 205 | c._send = Mock() 206 | c.PhysicalResourceId = True 207 | c._cfn_response(event) 208 | c._send.assert_called_once() 209 | self.assertEqual(True, c.PhysicalResourceId.startswith('test-stack-id_TestResourceId_')) 210 | 211 | c._send = Mock() 212 | c.PhysicalResourceId = '' 213 | event['PhysicalResourceId'] = 'pid-from-event' 214 | c._cfn_response(event) 215 | c._send.assert_called_once() 216 | self.assertEqual('pid-from-event', c.PhysicalResourceId) 217 | 218 | @patch('crhelper.log_helper.setupLogger', Mock()) 219 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 220 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 221 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 222 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 223 | def test_wrap_function(self): 224 | c = crhelper.resource_helper.CfnResource() 225 | 226 | def func(e, c): 227 | return 'testpid' 228 | 229 | c._wrap_function(func) 230 | self.assertEqual('testpid', c.PhysicalResourceId) 231 | self.assertNotEqual('FAILED', c.Status) 232 | 233 | def func(e, c): 234 | raise Exception('test exception') 235 | 236 | c._wrap_function(func) 237 | self.assertEqual('FAILED', c.Status) 238 | self.assertEqual('test exception', c.Reason) 239 | 240 | @patch('crhelper.log_helper.setupLogger', Mock()) 241 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 242 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 243 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 244 | def test_send(self): 245 | c = crhelper.resource_helper.CfnResource() 246 | s = Mock() 247 | c._send(send_response=s) 248 | s.assert_called_once() 249 | 250 | @patch('crhelper.log_helper.setupLogger', Mock()) 251 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 252 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 253 | @patch('crhelper.resource_helper.CfnResource._send', return_value=None) 254 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 255 | def test_timeout(self, s): 256 | c = crhelper.resource_helper.CfnResource() 257 | c._timeout() 258 | s.assert_called_with('FAILED', "Execution timed out") 259 | 260 | @patch('crhelper.log_helper.setupLogger', Mock()) 261 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 262 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 263 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 264 | def test_set_timeout(self): 265 | c = crhelper.resource_helper.CfnResource() 266 | c._context = MockContext() 267 | def func(): 268 | return None 269 | 270 | c._set_timeout() 271 | t = threading.Timer(1000, func) 272 | self.assertEqual(type(t), type(c._timer)) 273 | t.cancel() 274 | c._timer.cancel() 275 | 276 | @patch('crhelper.log_helper.setupLogger', Mock()) 277 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 278 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 279 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 280 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 281 | def test_cleanup_response(self): 282 | c = crhelper.resource_helper.CfnResource() 283 | c.Data = {"CrHelperPoll": 1, "CrHelperPermission": 2, "CrHelperRule": 3} 284 | c._cleanup_response() 285 | self.assertEqual({}, c.Data) 286 | 287 | @patch('crhelper.log_helper.setupLogger', Mock()) 288 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 289 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 290 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 291 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 292 | def test_remove_polling(self): 293 | c = crhelper.resource_helper.CfnResource() 294 | c._context = MockContext() 295 | 296 | c._events_client.remove_targets = Mock() 297 | c._events_client.delete_rule = Mock() 298 | c._lambda_client.remove_permission = Mock() 299 | 300 | with self.assertRaises(Exception) as e: 301 | c._remove_polling() 302 | 303 | self.assertEqual("failed to cleanup CloudWatch event polling", str(e)) 304 | c._events_client.remove_targets.assert_not_called() 305 | c._events_client.delete_rule.assert_not_called() 306 | c._lambda_client.remove_permission.assert_not_called() 307 | 308 | c._event["CrHelperRule"] = "1/2" 309 | c._event["CrHelperPermission"] = "1/2" 310 | c._remove_polling() 311 | c._events_client.remove_targets.assert_called() 312 | c._events_client.delete_rule.assert_called() 313 | c._lambda_client.remove_permission.assert_called() 314 | 315 | @patch('crhelper.log_helper.setupLogger', Mock()) 316 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 317 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 318 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 319 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 320 | @patch('crhelper.resource_helper.CfnResource._rand_string', Mock(return_value='PLURAL=1')) 321 | def test_setup_polling(self): 322 | c = crhelper.resource_helper.CfnResource() 323 | c._context = MockContext() 324 | c._event = test_events["Update"] 325 | c._lambda_client.add_permission = Mock() 326 | c._events_client.put_rule = Mock(return_value={"RuleArn": "arn:aws:lambda:blah:blah:function:blah/blah"}) 327 | c._events_client.put_targets = Mock() 328 | c._setup_polling() 329 | c._events_client.put_targets.assert_called() 330 | c._events_client.put_rule.assert_called_with(Name='TestResourceIdPLURAL=1', ScheduleExpression='rate(2 minutes)', State='ENABLED') 331 | c._lambda_client.add_permission.assert_called() 332 | 333 | @patch('crhelper.log_helper.setupLogger', Mock()) 334 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 335 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 336 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 337 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 338 | @patch('crhelper.resource_helper.CfnResource._rand_string', Mock(return_value='PLURAL=0')) 339 | def test_setup_polling_plural_0(self): 340 | c = crhelper.resource_helper.CfnResource(polling_interval=1) 341 | c._context = MockContext() 342 | c._event = test_events["Update"] 343 | c._lambda_client.add_permission = Mock() 344 | c._events_client.put_rule = Mock(return_value={"RuleArn": "arn:aws:lambda:blah:blah:function:blah/blah"}) 345 | c._events_client.put_targets = Mock() 346 | c._setup_polling() 347 | c._events_client.put_targets.assert_called() 348 | c._events_client.put_rule.assert_called_with(Name='TestResourceIdPLURAL=0', ScheduleExpression='rate(1 minute)', State='ENABLED') 349 | c._lambda_client.add_permission.assert_called() 350 | 351 | @patch('crhelper.log_helper.setupLogger', Mock()) 352 | @patch('crhelper.resource_helper.CfnResource._poll_enabled', Mock(return_value=False)) 353 | @patch('crhelper.resource_helper.CfnResource._wait_for_cwlogs', Mock()) 354 | @patch('crhelper.resource_helper.CfnResource._send', Mock()) 355 | @patch('crhelper.resource_helper.CfnResource._set_timeout', Mock()) 356 | def test_wrappers(self): 357 | c = crhelper.resource_helper.CfnResource() 358 | 359 | def func(): 360 | pass 361 | 362 | for f in ["create", "update", "delete", "poll_create", "poll_update", "poll_delete"]: 363 | self.assertEqual(None, getattr(c, "_%s_func" % f)) 364 | getattr(c, f)(func) 365 | self.assertEqual(func, getattr(c, "_%s_func" % f)) 366 | -------------------------------------------------------------------------------- /tests/test_utils.py: -------------------------------------------------------------------------------- 1 | import json 2 | from unittest.mock import patch, Mock, ANY, MagicMock 3 | from crhelper import utils 4 | import unittest 5 | import ssl 6 | import tempfile 7 | 8 | 9 | class TestLogHelper(unittest.TestCase): 10 | TEST_URL = "https://test_url/this/is/the/url?query=123#aaa" 11 | 12 | @patch('crhelper.utils.HTTPSConnection', autospec=True) 13 | def test_send_succeeded_response(self, https_connection_mock): 14 | utils._send_response(self.TEST_URL, {}) 15 | https_connection_mock.assert_called_once_with("test_url", context=ANY) 16 | https_connection_mock.return_value.request.assert_called_once_with( 17 | body='{}', 18 | headers={"content-type": "", "content-length": "2"}, 19 | method="PUT", 20 | url="/this/is/the/url?query=123#aaa", 21 | ) 22 | 23 | @patch('crhelper.utils.HTTPSConnection', autospec=True) 24 | def test_send_failed_response(self, https_connection_mock): 25 | utils._send_response(self.TEST_URL, Mock()) 26 | https_connection_mock.assert_called_once_with("test_url", context=ANY) 27 | response = json.loads(https_connection_mock.return_value.request.call_args[1]["body"]) 28 | expected_body = '{"Status": "FAILED", "Data": {}, "Reason": "' + response["Reason"] + '"}' 29 | https_connection_mock.return_value.request.assert_called_once_with( 30 | body=expected_body, 31 | headers={"content-type": "", "content-length": str(len(expected_body))}, 32 | method="PUT", 33 | url="/this/is/the/url?query=123#aaa", 34 | ) 35 | 36 | @patch('crhelper.utils.ssl.create_default_context', autospec=True) 37 | @patch('crhelper.utils.HTTPSConnection', autospec=True) 38 | def test_send_response_no_ssl_verify(self, https_connection_mock, ssl_create_context_mock): 39 | ctx_mock = Mock() 40 | ssl_create_context_mock.return_value = ctx_mock 41 | utils._send_response(self.TEST_URL, {}, ssl_verify=False) 42 | https_connection_mock.assert_called_once_with("test_url", context=ctx_mock) 43 | self.assertFalse(ctx_mock.check_hostname) 44 | self.assertEqual(ctx_mock.verify_mode, ssl.CERT_NONE) 45 | 46 | @patch('crhelper.utils.ssl.create_default_context', autospec=True) 47 | @patch('crhelper.utils.HTTPSConnection', autospec=True) 48 | def test_send_response_custom_ca(self, https_connection_mock, ssl_create_context_mock): 49 | ctx_mock = Mock() 50 | ssl_create_context_mock.return_value = ctx_mock 51 | with tempfile.NamedTemporaryFile() as tmp: 52 | utils._send_response(self.TEST_URL, {}, ssl_verify=tmp.name) 53 | https_connection_mock.assert_called_once_with("test_url", context=ctx_mock) 54 | ctx_mock.load_verify_locations.assert_called_once_with(cafile=tmp.name) 55 | 56 | @patch('crhelper.utils.ssl.create_default_context', autospec=True) 57 | @patch('crhelper.utils.HTTPSConnection', autospec=True) 58 | def test_send_response_non_existant_custom_ca(self, https_connection_mock, ssl_create_context_mock): 59 | ctx_mock = Mock() 60 | ssl_create_context_mock.return_value = ctx_mock 61 | utils._send_response(self.TEST_URL, {}, ssl_verify='/invalid/path/to/ca') 62 | https_connection_mock.assert_called_once_with("test_url", context=ANY) 63 | ctx_mock.load_verify_locations.assert_not_called() 64 | 65 | 66 | @patch('crhelper.utils.logger') 67 | @patch('crhelper.utils.HTTPSConnection') 68 | def test_send_response_retry(self, mock_https_connection, mock_logger): 69 | # Mock the behavior of HTTPSConnection to fail on the first two attempts 70 | mock_connection = mock_https_connection.return_value 71 | mock_connection.getresponse.side_effect = [ 72 | Exception("Unexpected failure sending response to CloudFormation"), 73 | Exception("Unexpected failure sending response to CloudFormation"), MagicMock(reason="OK")] 74 | 75 | # Call the code under test 76 | response_url = "https://example.com/response" 77 | response_body = {"Status": "SUCCESS", "Data": {"Message": "Test message"}} 78 | ssl_verify = True 79 | utils._send_response(response_url, response_body, ssl_verify) 80 | 81 | # Assert that the logger was called with the expected error and info messages 82 | self.assertEqual(mock_logger.error.call_count, 2) 83 | self.assertEqual(mock_logger.info.call_count, 1) 84 | 85 | # Assert that the HTTPSConnection was called the expected number of times 86 | self.assertEqual(mock_https_connection.call_count, 3) 87 | 88 | # Assert that the function retried the expected number of times 89 | self.assertEqual(mock_connection.getresponse.call_count, 3) -------------------------------------------------------------------------------- /tests/unit/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/aws-cloudformation/custom-resource-helper/b07bd4ec17e175a5148ee07e7d5efa12dd7bb6a5/tests/unit/__init__.py --------------------------------------------------------------------------------