Python disco_aws_automation.DiscoAWS类代码示例

    def test_provision_hc_with_chaos_using_config(self, mock_config, **kwargs):
        """
        Provision creates the proper launch configuration and autoscaling group with chaos from config
        """
        config_dict = get_default_config_dict()
        config_dict["mhcunittest"]["chaos"] = "True"
        aws = DiscoAWS(config=get_mock_config(config_dict), environment_name=TEST_ENV_NAME,
                       log_metrics=MagicMock())
        mock_ami = self._get_image_mock(aws)
        aws.update_elb = MagicMock(return_value=None)
        aws.discogroup.elastigroup.spotinst_client = MagicMock()

        with patch("disco_aws_automation.DiscoAWS.get_meta_network", return_value=_get_meta_network_mock()):
            with patch("boto.ec2.connection.EC2Connection.get_all_snapshots", return_value=[]):
                with patch("disco_aws_automation.DiscoAWS.create_scaling_schedule", return_value=None):
                    with patch("boto.ec2.autoscale.AutoScaleConnection.create_or_update_tags",
                               return_value=None):
                        metadata = aws.provision(ami=mock_ami, hostclass="mhcunittest",
                                                 owner="unittestuser",
                                                 min_size=1, desired_size=1, max_size=1)

        self.assertEqual(metadata["hostclass"], "mhcunittest")
        self.assertFalse(metadata["no_destroy"])
        self.assertTrue(metadata["chaos"])
        _lc = aws.discogroup.get_configs()[0]
        self.assertRegexpMatches(_lc.name, r".*_mhcunittest_[0-9]*")
        self.assertEqual(_lc.image_id, mock_ami.id)
        self.assertTrue(aws.discogroup.get_existing_group(hostclass="mhcunittest"))
        _ag = aws.discogroup.get_existing_groups()[0]
        self.assertRegexpMatches(_ag['name'], r"unittestenv_mhcunittest_[0-9]*")
        self.assertEqual(_ag['min_size'], 1)
        self.assertEqual(_ag['max_size'], 1)
        self.assertEqual(_ag['desired_capacity'], 1)

    def test_update_elb_all_defaults(self, mock_config, **kwargs):
        """
        update_elb calls get_or_create_elb with default port and protocol values if all are missing
        """
        aws = DiscoAWS(config=self._get_elb_config(), environment_name=TEST_ENV_NAME, elb=MagicMock())
        aws.elb.get_or_create_elb = MagicMock(return_value=MagicMock())
        aws.get_meta_network_by_name = _get_meta_network_mock()
        aws.elb.delete_elb = MagicMock()

        aws.update_elb("mhcelb", update_autoscaling=False)

        aws.elb.delete_elb.assert_not_called()
        aws.elb.get_or_create_elb.assert_called_once_with(
            'mhcelb',
            health_check_url='/foo',
            hosted_zone_name='example.com',
            port_config=DiscoELBPortConfig(
                [
                    DiscoELBPortMapping(80, 'HTTP', 80, 'HTTP'),
                ]
            ),
            security_groups=['sg-1234abcd'], elb_public=False,
            sticky_app_cookie=None, subnets=['s-1234abcd', 's-1234abcd', 's-1234abcd'],
            elb_dns_alias=None,
            connection_draining_timeout=300, idle_timeout=300, testing=False,
            tags={
                'environment': 'unittestenv',
                'hostclass': 'mhcelb',
                'is_testing': '0',
                'productline': 'mock_productline'
            },
            cross_zone_load_balancing=True,
            cert_name=None
        )

    def test_provision_hc_simple_with_no_chaos(self, mock_config, **kwargs):
        """
        Provision creates the proper launch configuration and autoscaling group with no chaos
        """
        aws = DiscoAWS(config=mock_config, environment_name=TEST_ENV_NAME)
        mock_ami = self._get_image_mock(aws)
        aws.log_metrics = MagicMock()
        aws.update_elb = MagicMock(return_value=None)

        with patch("disco_aws_automation.DiscoAWS.get_meta_network", return_value=_get_meta_network_mock()):
            with patch("boto.ec2.connection.EC2Connection.get_all_snapshots", return_value=[]):
                with patch("disco_aws_automation.DiscoAWS.create_scaling_schedule", return_value=None):
                    with patch("boto.ec2.autoscale.AutoScaleConnection.create_or_update_tags",
                               return_value=None):
                        metadata = aws.provision(ami=mock_ami, hostclass="mhcunittest",
                                                 owner="unittestuser",
                                                 min_size=1, desired_size=1, max_size=1,
                                                 chaos="False")

        self.assertEqual(metadata["hostclass"], "mhcunittest")
        self.assertFalse(metadata["no_destroy"])
        self.assertFalse(metadata["chaos"])
        _lc = aws.autoscale.get_configs()[0]
        self.assertRegexpMatches(_lc.name, r".*_mhcunittest_[0-9]*")
        self.assertEqual(_lc.image_id, mock_ami.id)
        self.assertTrue(aws.autoscale.has_group("mhcunittest"))
        _ag = aws.autoscale.get_groups()[0]
        self.assertEqual(_ag.name, "unittestenv_mhcunittest")
        self.assertEqual(_ag.min_size, 1)
        self.assertEqual(_ag.max_size, 1)
        self.assertEqual(_ag.desired_capacity, 1)

 def test_create_scaling_schedule_no_sched(self, mock_config, **kwargs):
     """test create_scaling_schedule with only desired schedule"""
     aws = DiscoAWS(config=mock_config, environment_name=TEST_ENV_NAME, discogroup=MagicMock())
     aws.create_scaling_schedule("1", "2", "5", hostclass="mhcboo")
     aws.discogroup.assert_has_calls([
         call.delete_all_recurring_group_actions(hostclass='mhcboo', group_name=None)
     ])

 def test_wait_for_autoscaling_using_amiid(self, mock_config, **kwargs):
     '''test wait for autoscaling using the ami id to identify the instances'''
     aws = DiscoAWS(config=mock_config, environment_name=TEST_ENV_NAME)
     instances = [{"InstanceId": "i-123123aa"}]
     aws.instances_from_amis = MagicMock(return_value=instances)
     aws.wait_for_autoscaling('ami-12345678', 1)
     aws.instances_from_amis.assert_called_with(['ami-12345678'], group_name=None, launch_time=None)

 def test_update_elb_delete(self, mock_config, **kwargs):
     '''Update ELB deletes ELBs that are no longer configured'''
     aws = DiscoAWS(config=mock_config, environment_name=TEST_ENV_NAME, elb=MagicMock())
     aws.elb.get_elb = MagicMock(return_value=True)
     aws.elb.delete_elb = MagicMock()
     aws.update_elb("mhcfoo", update_autoscaling=False)
     aws.elb.delete_elb.assert_called_once_with("mhcfoo")

def run():
    """Parses command line and dispatches the commands"""
    config = read_config()
    parser = get_parser()
    args = parser.parse_args()
    configure_logging(args.debug)

    environment_name = args.env or config.get("disco_aws", "default_environment")

    aws = DiscoAWS(config, environment_name=environment_name)
    if args.mode == "create":
        aws.disco_storage.create_ebs_snapshot(args.hostclass, args.size)
    elif args.mode == "list":
        for snapshot in aws.disco_storage.get_snapshots(args.hostclasses):
            print("{0:26} {1:13} {2:9} {3} {4:4}".format(
                snapshot.tags['hostclass'], snapshot.id, snapshot.status,
                snapshot.start_time, snapshot.volume_size))
    elif args.mode == "cleanup":
        aws.disco_storage.cleanup_ebs_snapshots(args.keep)
    elif args.mode == "capture":
        instances = instances_from_args(aws, args)
        if not instances:
            logging.warning("No instances found")
        for instance in instances:
            return_code, output = aws.remotecmd(
                instance, ["sudo /opt/wgen/bin/take_snapshot.sh"], user="snapshot")
            if return_code:
                raise Exception("Failed to snapshot instance {0}:\n {1}\n".format(instance, output))
            logging.info("Successfully snapshotted %s", instance)
    elif args.mode == "delete":
        for snapshot_id in args.snapshots:
            aws.disco_storage.delete_snapshot(snapshot_id)
    elif args.mode == "update":
        snapshot = aws.disco_storage.get_latest_snapshot(args.hostclass)
        aws.autoscale.update_snapshot(args.hostclass, snapshot.id, snapshot.volume_size)

    def test_create_userdata_without_spotinst(self, **kwargs):
        """
        create_userdata doesn't set 'spotinst' key
        """
        config_dict = get_default_config_dict()
        aws = DiscoAWS(config=get_mock_config(config_dict), environment_name=TEST_ENV_NAME)

        user_data = aws.create_userdata(hostclass="mhcunittest", owner="unittestuser", is_spotinst=False)
        self.assertEqual(user_data["is_spotinst"], "0")

 def test_instances_from_amis(self, mock_config, **kwargs):
     '''test get instances using ami ids '''
     aws = DiscoAWS(config=mock_config, environment_name=TEST_ENV_NAME)
     instance = create_autospec(boto.ec2.instance.Instance)
     instance.id = "i-123123aa"
     instances = [instance]
     aws.instances = MagicMock(return_value=instances)
     self.assertEqual(aws.instances_from_amis('ami-12345678'), instances)
     aws.instances.assert_called_with(filters={"image_id": 'ami-12345678'}, instance_ids=None)

    def test_create_userdata_with_zookeeper(self, **kwargs):
        """
        create_userdata sets 'zookeepers' key
        """
        config_dict = get_default_config_dict()
        aws = DiscoAWS(config=get_mock_config(config_dict), environment_name=TEST_ENV_NAME)

        user_data = aws.create_userdata(hostclass="mhcunittest", owner="unittestuser")
        self.assertEqual(user_data["zookeepers"], "[\\\"mhczookeeper-{}.example.com:2181\\\"]".format(
            aws.vpc.environment_name))

 def test_instances_from_amis_with_group_name(self, mock_config, **kwargs):
     '''test get instances using ami ids in a specified group name'''
     aws = DiscoAWS(config=mock_config, environment_name=TEST_ENV_NAME)
     instance = create_autospec(boto.ec2.instance.Instance)
     instance.id = "i-123123aa"
     instances = [instance]
     aws.instances_from_asgs = MagicMock(return_value=instances)
     aws.instances = MagicMock(return_value=instances)
     self.assertEqual(aws.instances_from_amis('ami-12345678', group_name='test_group'), instances)
     aws.instances_from_asgs.assert_called_with(['test_group'])

    def test_create_userdata_with_eip(self, **kwargs):
        """
        create_userdata sets 'eip' key when an EIP is required
        """
        config_dict = get_default_config_dict()
        eip = "54.201.250.76"
        config_dict["mhcunittest"]["eip"] = eip
        aws = DiscoAWS(config=get_mock_config(config_dict), environment_name=TEST_ENV_NAME)

        user_data = aws.create_userdata(hostclass="mhcunittest", owner="unittestuser")
        self.assertEqual(user_data["eip"], eip)

 def test_create_scaling_schedule_only_desired(self, mock_config, **kwargs):
     """test create_scaling_schedule with only desired schedule"""
     aws = DiscoAWS(config=mock_config, environment_name=TEST_ENV_NAME, discogroup=MagicMock())
     aws.create_scaling_schedule("1", "[email protected] 0 * * *:[email protected] 0 * * *", "5", hostclass="mhcboo")
     aws.discogroup.assert_has_calls([
         call.delete_all_recurring_group_actions(hostclass='mhcboo', group_name=None),
         call.create_recurring_group_action('1 0 * * *', hostclass='mhcboo', group_name=None,
                                            min_size=None, desired_capacity=2, max_size=None),
         call.create_recurring_group_action('6 0 * * *', hostclass='mhcboo', group_name=None,
                                            min_size=None, desired_capacity=3, max_size=None)
     ], any_order=True)

 def test_create_scaling_schedule_overlapping(self, mock_config, **kwargs):
     """test create_scaling_schedule with only desired schedule"""
     aws = DiscoAWS(config=mock_config, environment_name=TEST_ENV_NAME)
     aws.autoscale = MagicMock()
     aws.create_scaling_schedule("mhcboo",
                                 "[email protected] 0 * * *:[email protected] 0 * * *",
                                 "[email protected] 0 * * *:[email protected] 0 * * *",
                                 "[email protected] 0 * * *:[email protected] 0 * * *")
     aws.autoscale.assert_has_calls([
         call.delete_all_recurring_group_actions('mhcboo'),
         call.create_recurring_group_action('mhcboo', '1 0 * * *',
                                            min_size=1, desired_capacity=2, max_size=6),
         call.create_recurring_group_action('mhcboo', '6 0 * * *',
                                            min_size=2, desired_capacity=3, max_size=9)
     ], any_order=True)

    def test_provision_hostclass_sched_all_none(self, mock_config, **kwargs):
        """
        Provision creates the proper autoscaling group sizes with scheduled sizes
        """
        aws = DiscoAWS(config=mock_config, environment_name=TEST_ENV_NAME)
        aws.log_metrics = MagicMock()
        aws.update_elb = MagicMock(return_value=None)

        with patch("disco_aws_automation.DiscoAWS.get_meta_network", return_value=_get_meta_network_mock()):
            with patch("boto.ec2.connection.EC2Connection.get_all_snapshots", return_value=[]):
                with patch("disco_aws_automation.DiscoAWS.create_scaling_schedule", return_value=None):
                    with patch("boto.ec2.autoscale.AutoScaleConnection.create_or_update_tags",
                               return_value=None):
                        aws.provision(ami=self._get_image_mock(aws),
                                      hostclass="mhcunittest", owner="unittestuser",
                                      min_size="", desired_size="", max_size="")

        _ag0 = aws.autoscale.get_groups()[0]

        self.assertEqual(_ag0.min_size, 0)  # minimum of listed sizes
        self.assertEqual(_ag0.desired_capacity, 0)  # maximum of listed sizes
        self.assertEqual(_ag0.max_size, 0)  # maximum of listed sizes

        with patch("disco_aws_automation.DiscoAWS.get_meta_network", return_value=_get_meta_network_mock()):
            with patch("boto.ec2.connection.EC2Connection.get_all_snapshots", return_value=[]):
                with patch("disco_aws_automation.DiscoAWS.create_scaling_schedule", return_value=None):
                    with patch("boto.ec2.autoscale.AutoScaleConnection.create_or_update_tags",
                               return_value=None):
                        aws.provision(ami=self._get_image_mock(aws),
                                      hostclass="mhcunittest", owner="unittestuser",
                                      min_size="3", desired_size="6", max_size="9")

        _ag1 = aws.autoscale.get_groups()[0]

        self.assertEqual(_ag1.min_size, 3)  # minimum of listed sizes
        self.assertEqual(_ag1.desired_capacity, 6)  # maximum of listed sizes
        self.assertEqual(_ag1.max_size, 9)  # maximum of listed sizes

        with patch("disco_aws_automation.DiscoAWS.get_meta_network", return_value=_get_meta_network_mock()):
            with patch("boto.ec2.connection.EC2Connection.get_all_snapshots", return_value=[]):
                with patch("disco_aws_automation.DiscoAWS.create_scaling_schedule", return_value=None):
                    with patch("boto.ec2.autoscale.AutoScaleConnection.create_or_update_tags",
                               return_value=None):
                        aws.provision(ami=self._get_image_mock(aws),
                                      hostclass="mhcunittest", owner="unittestuser",
                                      min_size="", desired_size="", max_size="")

        _ag2 = aws.autoscale.get_groups()[0]

        self.assertEqual(_ag2.min_size, 3)  # minimum of listed sizes
        self.assertEqual(_ag2.desired_capacity, 6)  # maximum of listed sizes
        self.assertEqual(_ag2.max_size, 9)  # maximum of listed sizes