Test your AI Agents with the all-new Agent to Agent Testing Platform.Learn More

How to use update_usage_plan method in localstack

Best Python code snippet using localstack_python

test_apigw_usage_plan.py

Source:test_apigw_usage_plan.py

1#!/usr/bin/python2# TODO: License goes here3import library.apigw_usage_plan as apigw_usage_plan4from library.apigw_usage_plan import ApiGwUsagePlan5import mock6from mock import patch7from mock import create_autospec8from mock import ANY9import unittest10import boto11from botocore.exceptions import BotoCoreError12class TestApiGwUsagePlan(unittest.TestCase):13  def setUp(self):14    self.module = mock.MagicMock()15    self.module.check_mode = False16    self.module.exit_json = mock.MagicMock()17    self.module.fail_json = mock.MagicMock()18    self.usage_plan  = ApiGwUsagePlan(self.module)19    self.usage_plan.client = mock.MagicMock()20    self.usage_plan.module.params = {21      'name': 'testify',22      'description': 'test_description',23      'api_stages': [{'rest_api_id': 'id1', 'stage': 'stage1'}],24      'throttle_burst_limit': 111,25      'throttle_rate_limit': 222.0,26      'quota_limit': 333,27      'quota_offset': 444,28      'quota_period': 'WEEK',29      'state': 'present',30    }31    reload(apigw_usage_plan)32  def test_boto_module_not_found(self):33    # Setup Mock Import Function34    import __builtin__ as builtins35    real_import = builtins.__import__36    def mock_import(name, *args):37      if name == 'boto': raise ImportError38      return real_import(name, *args)39    with mock.patch('__builtin__.__import__', side_effect=mock_import):40      reload(apigw_usage_plan)41      ApiGwUsagePlan(self.module)42    self.module.fail_json.assert_called_with(msg='boto and boto3 are required for this module')43  def test_boto3_module_not_found(self):44    # Setup Mock Import Function45    import __builtin__ as builtins46    real_import = builtins.__import__47    def mock_import(name, *args):48      if name == 'boto3': raise ImportError49      return real_import(name, *args)50    with mock.patch('__builtin__.__import__', side_effect=mock_import):51      reload(apigw_usage_plan)52      ApiGwUsagePlan(self.module)53    self.module.fail_json.assert_called_with(msg='boto and boto3 are required for this module')54  @patch.object(apigw_usage_plan, 'boto3')55  def test_boto3_client_properly_instantiated(self, mock_boto):56    ApiGwUsagePlan(self.module)57    mock_boto.client.assert_called_once_with('apigateway')58  @patch.object(ApiGwUsagePlan, '_update_usage_plan', return_value=('hi', 'mom'))59  def test_process_request_calls_get_usage_plans_and_stores_result_when_invoked(self, m):60    resp = {61      'items': [62        {'name': 'not a match', 'id': 'usage_plan_id'},63        {'name': 'testify', 'id': 'usage_plan_id'},64      ],65    }66    self.usage_plan.client.get_usage_plans = mock.MagicMock(return_value=resp)67    self.usage_plan.process_request()68    self.assertEqual(resp['items'][1], self.usage_plan.me)69    self.usage_plan.client.get_usage_plans.assert_called_once_with()70  def test_process_request_stores_None_result_when_not_found_in_get_usage_plans_result(self):71    resp = {72      'items': [73        {'name': 'not a match', 'id': 'usage_plan_id'},74        {'name': 'also not a match', 'id': 'usage_plan_id'},75      ],76    }77    self.usage_plan.client.get_usage_plans = mock.MagicMock(return_value=resp)78    self.usage_plan.process_request()79    self.assertEqual(None, self.usage_plan.me)80    self.usage_plan.client.get_usage_plans.assert_called_once_with()81  def test_process_request_calls_fail_json_when_get_usage_plans_raises_exception(self):82    self.usage_plan.client.get_usage_plans = mock.MagicMock(side_effect=BotoCoreError())83    self.usage_plan.process_request()84    self.usage_plan.client.get_usage_plans.assert_called_once_with()85    self.usage_plan.module.fail_json.assert_called_once_with(86      msg='Error when getting usage_plans from boto3: An unspecified error occurred'87    )88  @patch.object(ApiGwUsagePlan, '_delete_usage_plan', return_value='Egah!')89  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value={'id': 'found'})90  def test_process_request_calls_exit_json_with_expected_value_after_successful_delete(self, mr, md):91    self.usage_plan.module.params = {92      'name': 'testify',93      'state': 'absent',94    }95    self.usage_plan.process_request()96    self.usage_plan.module.exit_json.assert_called_once_with(changed='Egah!', usage_plan=None)97  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan')98  def test_process_request_removes_api_stages_before_executing_delete(self, m):99    m.return_value = {100      'id': 'found',101      'apiStages': [{'apiId': '1', 'stage': 's1'}, {'apiId': '2', 'stage': 's2'}]102    }103    self.usage_plan.module.params = {104      'name': 'testify',105      'state': 'absent',106    }107    expected_patches = [108      {'op': 'remove', 'path': '/apiStages', 'value': '1:s1'},109      {'op': 'remove', 'path': '/apiStages', 'value': '2:s2'},110    ]111    self.usage_plan.process_request()112    self.usage_plan.client.update_usage_plan.assert_called_once_with(113      usagePlanId='found',114      patchOperations=mock.ANY115    )116    self.assertEqual(len(expected_patches), len(self.usage_plan.client.update_usage_plan.call_args[1]['patchOperations']))117    self.assertItemsEqual(expected_patches, self.usage_plan.client.update_usage_plan.call_args[1]['patchOperations'])118  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value={'id': 'found'})119  def test_process_request_calls_delete_usage_plan_when_state_absent_and_usage_plan_found(self, m):120    self.usage_plan.module.params = {121      'name': 'testify',122      'state': 'absent',123    }124    self.usage_plan.process_request()125    self.usage_plan.client.delete_usage_plan.assert_called_once_with(usagePlanId='found')126  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value={'id': 'found'})127  def test_process_request_skips_delete_and_calls_exit_json_with_true_when_check_mode_set_and_auth_found(self, m):128    self.usage_plan.module.params = {129      'name': 'testify',130      'state': 'absent',131    }132    self.usage_plan.module.check_mode = True133    self.usage_plan.process_request()134    self.assertEqual(0, self.usage_plan.client.delete_usage_plan.call_count)135    self.usage_plan.module.exit_json.assert_called_once_with(changed=True, usage_plan=None)136  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value={'id': 'found'})137  def test_process_request_calls_fail_json_when_delete_usage_plan_raises_error(self, m):138    self.usage_plan.module.params = {139      'name': 'testify',140      'state': 'absent',141    }142    self.usage_plan.client.delete_usage_plan = mock.MagicMock(side_effect=BotoCoreError)143    self.usage_plan.process_request()144    self.usage_plan.client.delete_usage_plan.assert_called_once_with(usagePlanId='found')145    self.usage_plan.module.fail_json.assert_called_once_with(146      msg='Error when deleting usage_plan via boto3: An unspecified error occurred'147    )148  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value=None)149  def test_process_request_skips_delete_when_usage_plan_not_found(self, m):150    self.usage_plan.module.params = {151      'name': 'testify',152      'state': 'absent',153    }154    self.usage_plan.process_request()155    self.assertEqual(0, self.usage_plan.client.delete_usage_plan.call_count)156  @patch.object(ApiGwUsagePlan, '_create_usage_plan', return_value=('eye', 'on the sandwich'))157  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value=None)158  def test_process_request_calls_exit_json_with_expected_value_after_successful_create(self, mra, mca):159    self.usage_plan.process_request()160    self.usage_plan.module.exit_json.assert_called_once_with(changed='eye', usage_plan='on the sandwich')161  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value=None)162  def test_process_request_returns_create_usage_plan_result_when_create_succeeds(self, m):163    self.usage_plan.client.create_usage_plan = mock.MagicMock(return_value='woot')164    self.usage_plan.process_request()165    self.usage_plan.module.exit_json.assert_called_once_with(changed=True, usage_plan='woot')166  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value=None)167  def test_process_request_calls_create_usage_plan_when_state_present_and_usage_plan_not_found(self, m):168    self.usage_plan.process_request()169    self.usage_plan.client.create_usage_plan.assert_called_once_with(170      name='testify',171      description='test_description',172      apiStages=[{'apiId': 'id1', 'stage': 'stage1'}],173      throttle={'burstLimit': 111, 'rateLimit': 222},174      quota={'limit': 333, 'offset': 444, 'period': 'WEEK'},175    )176  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value=None)177  def test_process_request_calls_fail_json_when_create_usage_plan_raises_exception(self, m):178    self.usage_plan.module.params = {179      'name': 'testify',180      'state': 'present',181    }182    self.usage_plan.client.create_usage_plan = mock.MagicMock(side_effect=BotoCoreError())183    self.usage_plan.process_request()184    self.usage_plan.client.create_usage_plan.assert_called_once_with(name='testify')185    self.usage_plan.module.fail_json.assert_called_once_with(186      msg='Error when creating usage_plan via boto3: An unspecified error occurred'187    )188  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value=None)189  def test_process_request_skips_create_call_and_returns_changed_True_when_check_mode(self, m):190    self.usage_plan.module.check_mode = True191    self.usage_plan.process_request()192    self.assertEqual(0, self.usage_plan.client.create_usage_plan.call_count)193    self.usage_plan.module.exit_json.assert_called_once_with(changed=True, usage_plan=None)194  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value={'id': 'hi'})195  def test_process_request_calls_update_usage_plan_for_replace_ops_when_state_present_and_usage_plan_changed(self, m):196    m.return_value = {197      'name': 'testify',198      'id': 'ab12345',199      'description': 'old and busted',200      'apiStages': [{'apiId': 'id1', 'stage': 'stage1'}],201      'throttle': {'burstLimit': 'old', 'rateLimit': 'old'},202      'quota': {'limit': 'old', 'offset': 'old', 'period': 'old'},203    }204    expected_patches = [205      {'op': 'replace', 'path': '/description', 'value': 'test_description'},206      {'op': 'replace', 'path': '/quota/period', 'value': 'WEEK'},207      {'op': 'replace', 'path': '/quota/offset', 'value': '444'},208      {'op': 'replace', 'path': '/quota/limit', 'value': '333'},209      {'op': 'replace', 'path': '/throttle/rateLimit', 'value': '222.0'},210      {'op': 'replace', 'path': '/throttle/burstLimit', 'value': '111'},211    ]212    self.usage_plan.process_request()213    self.usage_plan.client.update_usage_plan.assert_called_once_with(214      usagePlanId='ab12345',215      patchOperations=mock.ANY216    )217    self.assertEqual(len(expected_patches), len(self.usage_plan.client.update_usage_plan.call_args[1]['patchOperations']))218    self.assertItemsEqual(expected_patches, self.usage_plan.client.update_usage_plan.call_args[1]['patchOperations'])219  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value={'id': 'hi'})220  def test_process_request_calls_update_usage_plan_for_remove_ops_when_state_present_and_usage_plan_changed(self, m):221    self.usage_plan.module.params = {222      'name': 'testify',223      'description': 'test_description',224      'api_stages': [],225      'state': 'present',226    }227    m.return_value = {228      'name': 'testify',229      'id': 'ab12345',230      'description': 'test_description',231      'apiStages': [{'apiId': 'id1', 'stage': 'stage1'},{'apiId': 'id2', 'stage': 'stage2'}],232      'throttle': {'burstLimit': 'old', 'rateLimit': 'old'},233      'quota': {'limit': 'old', 'offset': 'old', 'period': 'old'},234    }235    expected_patches = [236      {'op': 'remove', 'path': '/quota'},237      {'op': 'remove', 'path': '/throttle'},238      {'op': 'remove', 'path': '/apiStages', 'value': 'id1:stage1'},239      {'op': 'remove', 'path': '/apiStages', 'value': 'id2:stage2'},240    ]241    self.usage_plan.process_request()242    self.usage_plan.client.update_usage_plan.assert_called_once_with(243      usagePlanId='ab12345',244      patchOperations=mock.ANY245    )246    self.assertEqual(len(expected_patches), len(self.usage_plan.client.update_usage_plan.call_args[1]['patchOperations']))247    self.assertItemsEqual(expected_patches, self.usage_plan.client.update_usage_plan.call_args[1]['patchOperations'])248  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value={'id': 'hi'})249  def test_process_request_calls_update_usage_plan_for_add_ops_when_state_present_and_usage_plan_changed(self, m):250    self.usage_plan.module.params = {251      'name': 'testify',252      'api_stages': [{'rest_api_id': 'id1', 'stage': 'stage1'}],253      'throttle_burst_limit': 111,254      'throttle_rate_limit': 222.0,255      'quota_limit': 333,256      'quota_offset': 444,257      'quota_period': 'WEEK',258    }259    m.return_value = {260      'name': 'testify',261      'id': 'ab12345',262      'description': 'old and busted',263      'apiStages': [{'apiId': 'idx', 'stage': 'stagex'}],264    }265    expected_patches = [266      {'op': 'replace', 'path': '/description', 'value': ''},267      {'op': 'add', 'path': '/quota/period', 'value': 'WEEK'},268      {'op': 'add', 'path': '/quota/offset', 'value': '444'},269      {'op': 'add', 'path': '/quota/limit', 'value': '333'},270      {'op': 'add', 'path': '/throttle/rateLimit', 'value': '222.0'},271      {'op': 'add', 'path': '/throttle/burstLimit', 'value': '111'},272      {'op': 'add', 'path': '/apiStages', 'value': 'id1:stage1'},273    ]274    self.usage_plan.process_request()275    self.usage_plan.client.update_usage_plan.assert_called_once_with(276      usagePlanId='ab12345',277      patchOperations=mock.ANY278    )279    self.assertEqual(len(expected_patches), len(self.usage_plan.client.update_usage_plan.call_args[1]['patchOperations']))280    self.assertItemsEqual(expected_patches, self.usage_plan.client.update_usage_plan.call_args[1]['patchOperations'])281  @patch('library.apigw_usage_plan.ApiGwUsagePlan._create_patches', return_value=[])282  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value={'something': 'here'})283  def test_process_request_skips_update_usage_plan_and_replies_false_when_no_changes(self, m, mcp):284    self.usage_plan.process_request()285    self.assertEqual(0, self.usage_plan.client.update_usage_plan.call_count)286    self.usage_plan.module.exit_json.assert_called_once_with(changed=False, usage_plan={'something': 'here'})287  @patch('library.apigw_usage_plan.ApiGwUsagePlan._create_patches', return_value=['patches!'])288  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value={'id': 'hi'})289  def test_process_request_calls_fail_json_when_update_usage_plan_raises_exception(self, m, mcp):290    self.usage_plan.client.update_usage_plan = mock.MagicMock(side_effect=BotoCoreError())291    self.usage_plan.process_request()292    self.usage_plan.client.update_usage_plan.assert_called_once_with(293      usagePlanId='hi',294      patchOperations=['patches!']295    )296    self.usage_plan.module.fail_json.assert_called_once_with(297      msg='Error when updating usage_plan via boto3: An unspecified error occurred'298    )299  @patch('library.apigw_usage_plan.ApiGwUsagePlan._create_patches', return_value=['patches!'])300  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', side_effect=[{'id': 'hi'}, 'second call'])301  def test_process_request_returns_result_of_find_when_update_is_successful(self, m, mcp):302    self.usage_plan.process_request()303    self.usage_plan.client.update_usage_plan.assert_called_once_with(304      usagePlanId='hi',305      patchOperations=['patches!']306    )307    self.usage_plan.module.exit_json.assert_called_once_with(changed=True, usage_plan='second call')308  @patch('library.apigw_usage_plan.ApiGwUsagePlan._create_patches', return_value=['patches!'])309  @patch.object(ApiGwUsagePlan, '_retrieve_usage_plan', return_value={'something': 'here'})310  def test_process_request_skips_update_usage_plan_and_replies_true_when_check_mode(self, m, mcp):311    self.usage_plan.module.check_mode = True312    self.usage_plan.process_request()313    self.assertEqual(0, self.usage_plan.client.update_usage_plan.call_count)314    self.usage_plan.module.exit_json.assert_called_once_with(changed=True, usage_plan={'something': 'here'})315  def test_define_argument_spec(self):316    result = ApiGwUsagePlan._define_module_argument_spec()317    self.assertIsInstance(result, dict)318    self.assertEqual(result, dict(319                     name=dict(required=True),320                     description=dict(required=False, default=''),321                     api_stages=dict(322                       type='list',323                       required=False,324                       default=[],325                       rest_api_id=dict(required=True),326                       stage=dict(required=True)327                     ),328                     throttle_burst_limit=dict(required=False, default=-1, type='int'),329                     throttle_rate_limit=dict(required=False, default=-1.0, type='float'),330                     quota_limit=dict(required=False, default=-1, type='int'),331                     quota_offset=dict(required=False, default=-1, type='int'),332                     quota_period=dict(required=False, default='', choices=['', 'DAY','WEEK','MONTH']),333                     state=dict(default='present', choices=['present', 'absent']),334    ))335  @patch.object(apigw_usage_plan, 'AnsibleModule')336  @patch.object(apigw_usage_plan, 'ApiGwUsagePlan')337  def test_main(self, mock_ApiGwUsagePlan, mock_AnsibleModule):338    mock_ApiGwUsagePlan_instance      = mock.MagicMock()339    mock_AnsibleModule_instance     = mock.MagicMock()340    mock_ApiGwUsagePlan.return_value  = mock_ApiGwUsagePlan_instance341    mock_AnsibleModule.return_value = mock_AnsibleModule_instance342    apigw_usage_plan.main()343    mock_ApiGwUsagePlan.assert_called_once_with(mock_AnsibleModule_instance)344    assert mock_ApiGwUsagePlan_instance.process_request.call_count == 1345if __name__ == '__main__':...

update_usage_plan.py

Source:update_usage_plan.py

1# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.2# SPDX-License-Identifier: MIT-03import json4import boto35import logger6from crhelper import CfnResource7helper = CfnResource()8try:9    dynamodb = boto3.resource('dynamodb')10    apigateway = boto3.client('apigateway')11except Exception as e:12    helper.init_failure(e)13@helper.create14def do_action(event, _):15    """ Usage plans are created as part of bootstrap template.16        This method associates the usage plans for various tiers with tenant Apis17    Args:18        event ([type]): [description]19        _ ([type]): [description]20    """21    logger.info("adding api gateway stage to usage plan")22    api_id = event['ResourceProperties']['ApiGatewayId']23    settings_table_name = event['ResourceProperties']['SettingsTableName']24    is_pooled_deploy = event['ResourceProperties']['IsPooledDeploy']25    stage = event['ResourceProperties']['Stage']26    usage_plan_id_basic = event['ResourceProperties']['UsagePlanBasicTier']27    usage_plan_id_standard = event['ResourceProperties']['UsagePlanStandardTier']28    usage_plan_id_premium = event['ResourceProperties']['UsagePlanPremiumTier']29    usage_plan_id_platinum = event['ResourceProperties']['UsagePlanPlatinumTier']30    table_system_settings = dynamodb.Table(settings_table_name)31    if(is_pooled_deploy == "true"):32        response_apigateway = apigateway.update_usage_plan (33                usagePlanId=usage_plan_id_basic,34                patchOperations=[35                    {36                        'op':'add',37                        'path':'/apiStages',38                        'value': api_id + ":" + stage39                    }40                ]41        )42        response_apigateway = apigateway.update_usage_plan (43                usagePlanId=usage_plan_id_standard,44                patchOperations=[45                    {46                        'op':'add',47                        'path':'/apiStages',48                        'value': api_id + ":" + stage49                    }50                ]51        )52        53        response_apigateway = apigateway.update_usage_plan (54                usagePlanId=usage_plan_id_premium,55                patchOperations=[56                    {57                        'op':'add',58                        'path':'/apiStages',59                        'value': api_id + ":" + stage60                    }61                ]62        )63        64    else:65        66        response_apigateway = apigateway.update_usage_plan (67                usagePlanId=usage_plan_id_platinum,68                patchOperations=[69                    {70                        'op':'add',71                        'path':'/apiStages',72                        'value': api_id + ":" + stage73                    }74                ]75        )76        77@helper.update78@helper.delete79def do_nothing(_, __):80    pass81def handler(event, context):   82    helper(event, context)83        ...

Automation Testing Tutorials

Learn to execute automation testing from scratch with LambdaTest Learning Hub. Right from setting up the prerequisites to run your first automation test, to following best practices and diving deeper into advanced test scenarios. LambdaTest Learning Hubs compile a list of step-by-step guides to help you be proficient with different test automation frameworks i.e. Selenium, Cypress, TestNG etc.