from unittest import TestCase

from samtranslator.model.update_policy import UpdatePolicy


class TestUpdatePolicy(TestCase):
    def test_to_dict_only_application_and_deployment_group(self):
        expected_dict = {
            'CodeDeployLambdaAliasUpdate': {
                'ApplicationName': 'application_name',
                'DeploymentGroupName': 'deployment_group_name'
            }
        }
        update_policy_dict = UpdatePolicy('application_name', 'deployment_group_name', None, None).to_dict()

        self.assertEqual(expected_dict, update_policy_dict)

    def test_to_dict_all_fields(self):
        expected_dict = {
            'CodeDeployLambdaAliasUpdate': {
                'ApplicationName': 'application_name',
                'DeploymentGroupName': 'deployment_group_name',
                'BeforeAllowTrafficHook': 'before_allow_traffic_hook',
                'AfterAllowTrafficHook': 'after_allow_traffic_hook'
            }
        }
        update_policy_dict = UpdatePolicy('application_name',
                                          'deployment_group_name',
                                          'before_allow_traffic_hook',
                                          'after_allow_traffic_hook').to_dict()

        self.assertEqual(expected_dict, update_policy_dict)
