<commit_before>from gevent import monkey; monkey.patch_all()
from corehq.elastic import get_es_new

from pillowtop.es_utils import pillow_index_exists, get_all_elasticsearch_pillow_classes, \
    get_all_expected_es_indices, needs_reindex

from cStringIO import StringIO
import traceback
from datetime import datetime
from django.core.mail import mail_admins
from corehq.pillows.user import add_demo_user_to_user_index
import gevent
from django.core.management.base import BaseCommand
from django.core.management import call_command
from django.conf import settings

def get_reindex_commands(pillow_class_name):
    # pillow_command_map is a mapping from es pillows
    # to lists of management commands or functions
    # that should be used to rebuild the index from scratch
    pillow_command_map = {
        'DomainPillow': ['ptop_fast_reindex_domains'],
        'CasePillow': ['ptop_fast_reindex_cases'],
        'XFormPillow': ['ptop_fast_reindex_xforms'],
        # groupstousers indexing must happen after all users are indexed
        'UserPillow': [
            'ptop_fast_reindex_users',
            add_demo_user_to_user_index,
            'ptop_fast_reindex_groupstousers',
            # 'ptop_fast_reindex_unknownusers',  removed until we have a better workflow for this
        ],
        'AppPillow': ['ptop_fast_reindex_apps'],
        'GroupPillow': ['ptop_fast_reindex_groups'],
        'ReportXFormPillow': ['ptop_fast_reindex_reportxforms'],
        'ReportCasePillow': ['ptop_fast_reindex_reportcases'],
    }
    return pillow_command_map.get(pillow_class_name, [])

def do_reindex(pillow_class_name):
    print "Starting pillow preindex %s" % pillow_class_name
    reindex_commands = get_reindex_commands(pillow_class_name)
    for reindex_command in reindex_commands:
        if isinstance(reindex_command, basestring):
            call_command(reindex_command, **{'noinput': True, 'bulk': True})
        else:
            reindex_command()
    print "Pillow preindex finished %s" % pillow_class_name

class Command(BaseCommand):
    help = ("Preindex ES pillows. "
            "Only run reindexer if the index doesn't exist.")

    def handle(self, *args, **options):
        runs = []
        aliased_classes = get_all_elasticsearch_pillow_classes()
        all_es_indices = get_all_expected_es_indices()
        es = get_es_new()
        indices_needing_reindex = [info for info in all_es_indices if needs_reindex(es, info)]
        aliasable_pillows = [p(online=False) for p in aliased_classes]

        print "Indices needing reindex:"
        unmapped_indices = [info.index for info in indices_needing_reindex]
        print unmapped_indices

        print ("Reindexing master pillows that do not exist yet "
               "(ones with aliases skipped)")

        pillows_without_index = [pillow for pillow in aliasable_pillows
                                 if not pillow_index_exists(pillow)]

        reindex_pillows = filter(lambda x: x.es_index in unmapped_indices,
                                 pillows_without_index)

        print "Reindexing:\n\t",
        print '\n\t'.join(x.es_index for x in reindex_pillows)

        if len(reindex_pillows) > 0:
            preindex_message = """
        Heads up!

        %s is going to start preindexing the following pillows:
        %s

        This may take a while, so don't deploy until all these have reported finishing.
            """ % (
                settings.EMAIL_SUBJECT_PREFIX,
                ', '.join([x.__class__.__name__ for x in reindex_pillows])
            )

            mail_admins("Pillow preindexing starting", preindex_message)

        start = datetime.utcnow()
        for pillow in reindex_pillows:
            # loop through pillows once before running greenlets
            # to fail hard on misconfigured pillows
            pillow_class_name = pillow.__class__.__name__
            reindex_command = get_reindex_commands(pillow_class_name)
            if not reindex_command:
                raise Exception(
                    "Error, pillow [%s] is not configured "
                    "with its own management command reindex command "
                    "- it needs one" % pillow_class_name
                )

        for pillow in reindex_pillows:
            print pillow.__class__.__name__
            g = gevent.spawn(do_reindex, pillow.__class__.__name__)
            runs.append(g)

        if len(reindex_pillows) > 0:
            gevent.joinall(runs)
            try:
                for job in runs:
                    job.get()
            except Exception:
                f = StringIO()
                traceback.print_exc(file=f)
                mail_admins("Pillow preindexing failed", f.getvalue())
                raise
            else:
                mail_admins(
                    "Pillow preindexing completed",
                    "Reindexing %s took %s seconds" % (
                        ', '.join([x.__class__.__name__ for x in reindex_pillows]),
                        (datetime.utcnow() - start).seconds
                    )
                )

        print "All pillowtop reindexing jobs completed"
<commit_msg>simplify logic around preindex workflow<commit_after>from gevent import monkey; monkey.patch_all()
from corehq.elastic import get_es_new

from pillowtop.es_utils import get_all_elasticsearch_pillow_classes, get_all_expected_es_indices

from cStringIO import StringIO
import traceback
from datetime import datetime
from django.core.mail import mail_admins
from corehq.pillows.user import add_demo_user_to_user_index
import gevent
from django.core.management.base import BaseCommand
from django.core.management import call_command
from django.conf import settings

def get_reindex_commands(pillow_class_name):
    # pillow_command_map is a mapping from es pillows
    # to lists of management commands or functions
    # that should be used to rebuild the index from scratch
    pillow_command_map = {
        'DomainPillow': ['ptop_fast_reindex_domains'],
        'CasePillow': ['ptop_fast_reindex_cases'],
        'XFormPillow': ['ptop_fast_reindex_xforms'],
        # groupstousers indexing must happen after all users are indexed
        'UserPillow': [
            'ptop_fast_reindex_users',
            add_demo_user_to_user_index,
            'ptop_fast_reindex_groupstousers',
            # 'ptop_fast_reindex_unknownusers',  removed until we have a better workflow for this
        ],
        'AppPillow': ['ptop_fast_reindex_apps'],
        'GroupPillow': ['ptop_fast_reindex_groups'],
        'ReportXFormPillow': ['ptop_fast_reindex_reportxforms'],
        'ReportCasePillow': ['ptop_fast_reindex_reportcases'],
    }
    return pillow_command_map.get(pillow_class_name, [])

def do_reindex(pillow_class_name):
    print "Starting pillow preindex %s" % pillow_class_name
    reindex_commands = get_reindex_commands(pillow_class_name)
    for reindex_command in reindex_commands:
        if isinstance(reindex_command, basestring):
            call_command(reindex_command, **{'noinput': True, 'bulk': True})
        else:
            reindex_command()
    print "Pillow preindex finished %s" % pillow_class_name

class Command(BaseCommand):
    help = ("Preindex ES pillows. "
            "Only run reindexer if the index doesn't exist.")

    def handle(self, *args, **options):
        runs = []
        aliased_classes = get_all_elasticsearch_pillow_classes()
        all_es_indices = get_all_expected_es_indices()
        es = get_es_new()
        indices_needing_reindex = [info for info in all_es_indices if not es.indices.exists(info.index)]

        if not indices_needing_reindex:
            print 'Nothing needs to be reindexed'
            return

        aliasable_pillows = [p(online=False) for p in aliased_classes]
        index_names_needing_reindex = [info.index for info in indices_needing_reindex]

        print "Reindexing:\n\t",
        print '\n\t'.join(index_names_needing_reindex)

        reindex_pillows = filter(lambda x: x.es_index in index_names_needing_reindex, aliasable_pillows)
        preindex_message = """
        Heads up!

        %s is going to start preindexing the following pillows:
        %s

        This may take a while, so don't deploy until all these have reported finishing.
            """ % (
                settings.EMAIL_SUBJECT_PREFIX,
                ', '.join([x.__class__.__name__ for x in reindex_pillows])
            )

        mail_admins("Pillow preindexing starting", preindex_message)

        start = datetime.utcnow()
        for pillow in reindex_pillows:
            # loop through pillows once before running greenlets
            # to fail hard on misconfigured pillows
            pillow_class_name = pillow.__class__.__name__
            reindex_command = get_reindex_commands(pillow_class_name)
            if not reindex_command:
                raise Exception(
                    "Error, pillow [%s] is not configured "
                    "with its own management command reindex command "
                    "- it needs one" % pillow_class_name
                )

        for pillow in reindex_pillows:
            print pillow.__class__.__name__
            g = gevent.spawn(do_reindex, pillow.__class__.__name__)
            runs.append(g)

        if len(reindex_pillows) > 0:
            gevent.joinall(runs)
            try:
                for job in runs:
                    job.get()
            except Exception:
                f = StringIO()
                traceback.print_exc(file=f)
                mail_admins("Pillow preindexing failed", f.getvalue())
                raise
            else:
                mail_admins(
                    "Pillow preindexing completed",
                    "Reindexing %s took %s seconds" % (
                        ', '.join([x.__class__.__name__ for x in reindex_pillows]),
                        (datetime.utcnow() - start).seconds
                    )
                )

        print "All pillowtop reindexing jobs completed"
<|endoftext|>


<commit_before>import pkg_resources
from os.path import join

import yaml
import pytest
from jinja2 import Environment, FileSystemLoader

from cosmo_tester.framework.util import (generate_ca_cert,
                                         generate_ssl_certificate,
                                         get_resource_path)

RESOURCES_PATH = pkg_resources.resource_filename(
    'cosmo_tester', 'test_suites/cluster/cfy_cluster_manager_resources')
REMOTE_SSH_KEY_PATH = '/tmp/cfy_cluster_manager_ssh_key.pem'
REMOTE_LICENSE_PATH = '/tmp/cfy_cluster_manager_license.yaml'
REMOTE_CERTS_PATH = '/tmp/certs'
REMOTE_CONFIGS_PATH = '/tmp/config_files'

@pytest.fixture()
def basic_config_dict(ssh_key, test_config):
    return {
        'ssh_key_path': REMOTE_SSH_KEY_PATH,
        'ssh_user': 'centos',
        'cloudify_license_path': REMOTE_LICENSE_PATH,
        'manager_rpm_download_link': test_config['cfy_cluster_manager'][
            'manager_install_rpm_path']
    }

@pytest.fixture()
def three_nodes_config_dict(basic_config_dict):
    return _get_config_dict('three_nodes_config.yaml', basic_config_dict)

@pytest.fixture()
def nine_nodes_config_dict(basic_config_dict):
    return _get_config_dict('nine_nodes_config.yaml', basic_config_dict)

@pytest.fixture()
def local_certs_path(tmp_path):
    dir_path = tmp_path / 'certs'
    dir_path.mkdir()
    return dir_path

@pytest.fixture()
def local_config_files(tmp_path):
    dir_path = tmp_path / 'config_files'
    dir_path.mkdir()
    return dir_path

def _get_config_dict(config_file_name, basic_config_dict):
    config_path = join(RESOURCES_PATH, config_file_name)
    with open(config_path) as config_file:
        config_dict = yaml.safe_load(config_file)

    config_dict.update(basic_config_dict)
    return config_dict

def test_create_three_nodes_cluster(three_vms, three_nodes_config_dict,
                                    test_config, ssh_key):
    """Tests that a three nodes cluster is successfully created."""
    node1, node2, node3 = three_vms
    _update_three_nodes_config_dict_vms(three_nodes_config_dict,
                                        [node1, node2, node3])

    _install_cluster(node1, three_nodes_config_dict, test_config, ssh_key)

def test_create_nine_nodes_cluster(nine_vms, nine_nodes_config_dict,
                                   test_config, ssh_key):
    """Tests that a nine nodes cluster is successfully created."""
    node1, node2, node3, node4, node5, node6, node7, node8, node9 = nine_vms
    for i, node in enumerate([node1, node2, node3, node4, node5, node6,
                              node7, node8, node9]):
        node_num = (i % 3) + 1
        if i < 3:
            node_name = 'manager-{0}'.format(node_num)
        elif i < 6:
            node_name = 'rabbitmq-{0}'.format(node_num)
        else:
            node_name = 'postgresql-{0}'.format(node_num)

        nine_nodes_config_dict['existing_vms'][node_name].update({
            'private_ip': str(node.private_ip_address),
            'public_ip': str(node.ip_address)
        })

    _install_cluster(node1, nine_nodes_config_dict, test_config, ssh_key)

def test_create_three_nodes_cluster_using_certificates(
        three_vms, three_nodes_config_dict, test_config,
        ssh_key, local_certs_path, logger):
    """Tests that the supllied certificates are being used in the cluster."""
    node1, node2, node3 = three_vms
    nodes_list = [node1, node2, node3]

    logger.info('Creating certificates')
    _create_certificates(local_certs_path, nodes_list)

    logger.info('Copying certificates to node-1')
    node1.run_command('mkdir -p {0}'.format(REMOTE_CERTS_PATH))
    for cert in local_certs_path.iterdir():
        node1.put_remote_file(local_path=str(cert),
                              remote_path=join(REMOTE_CERTS_PATH, cert.name))

    logger.info('Preparing cluster install configuration file')
    _update_three_nodes_config_dict_vms(three_nodes_config_dict, nodes_list)
    three_nodes_config_dict['ca_cert_path'] = join(REMOTE_CERTS_PATH, 'ca.pem')
    for i, node in enumerate(nodes_list, start=1):
        three_nodes_config_dict['existing_vms']['node-{0}'.format(i)].update({
            'cert_path': join(REMOTE_CERTS_PATH, 'node-{0}.crt'.format(i)),
            'key_path': join(REMOTE_CERTS_PATH, 'node-{0}.key'.format(i))
        })

    logger.info('Installing cluster')
    _install_cluster(node1, three_nodes_config_dict, test_config, ssh_key)

    logger.info('Asserting certs were successfully copied')
    local_ca_path = local_certs_path / 'ca.pem'
    ca_path_in_use = '/etc/cloudify/ssl/cloudify_internal_ca_cert.pem'
    for i, node in enumerate(nodes_list, start=1):
        node_name = 'node-{0}'.format(i)
        logger.info('Asserting certificates for %s', node_name)
        local_node_cert_path = local_certs_path / '{0}.crt'.format(node_name)
        local_node_key_path = local_certs_path / '{0}.key'.format(node_name)
        cert_path_in_use = '/etc/cloudify/ssl/cloudify_internal_cert.pem'
        key_path_in_use = '/etc/cloudify/ssl/cloudify_internal_key.pem'

        assert (local_node_cert_path.read_text() ==
                node.get_remote_file_content(cert_path_in_use))

        assert (local_node_key_path.read_text() ==
                node.get_remote_file_content(key_path_in_use))

        assert (local_ca_path.read_text() ==
                node.get_remote_file_content(ca_path_in_use))

def test_three_nodes_using_provided_config_files(
        three_vms, three_nodes_config_dict, test_config,
        ssh_key, local_certs_path, local_config_files, logger):
    node1, node2, node3 = three_vms
    nodes_list = [node1, node2, node3]
    logger.info('Creating certificates and passing them to the instances')
    node1.run_command('mkdir -p {0}'.format(REMOTE_CERTS_PATH))
    _create_certificates(local_certs_path, nodes_list, pass_certs=True)

    logger.info('Preparing config files')
    _prepare_config_files(nodes_list, local_config_files)
    _update_three_nodes_config_dict_vms(three_nodes_config_dict, nodes_list)
    for i, node in enumerate(nodes_list, start=1):
        three_nodes_config_dict['existing_vms']['node-{0}'.format(i)][
            'config_path'].update({
                'manager_config_path': node.remote_manager_config_path,
                'postgresql_config_path': node.remote_postgresql_config_path,
                'rabbitmq_config_path': node.remote_rabbitmq_config_path
            })

    logger.info('Installing cluster')
    _install_cluster(node1, three_nodes_config_dict, test_config, ssh_key)

    logger.info('Asserting config_files')
    base_cfy_dir = '/etc/cloudify'
    for i, node in enumerate(nodes_list, start=1):
        logger.info('Asserting config.yaml files for %s', 'node-{0}'.format(i))

        assert (node.local_manager_config_path.read_text() ==
                node.get_remote_file_content(
                    join(base_cfy_dir, node.local_manager_config_path.name))
                )

        assert (node.local_postgresql_config_path.read_text() ==
                node.get_remote_file_content(
                    join(base_cfy_dir, node.local_postgresql_config_path.name))
                )

        assert (node.local_rabbitmq_config_path.read_text() ==
                node.get_remote_file_content(
                    join(base_cfy_dir, node.local_rabbitmq_config_path.name))
                )

def _prepare_config_files(nodes_list, local_config_files):
    rabbitmq_cluster = {
        node.hostname: {
            'networks': {
                'default': str(node.private_ip_address)
            }
        } for node in nodes_list
    }

    postgresql_cluster = {
        node.hostname: {
            'ip': str(node.private_ip_address)
        } for node in nodes_list
    }

    nodes_list[0].run_command('mkdir -p {0}'.format(REMOTE_CONFIGS_PATH))
    templates_env = Environment(loader=FileSystemLoader(
        join(RESOURCES_PATH, 'config_files_templates')))
    _prepare_manager_config_files(
        templates_env.get_template('manager_config.yaml'),
        nodes_list, rabbitmq_cluster, postgresql_cluster, local_config_files)
    _prepare_postgresql_config_files(
        templates_env.get_template('postgresql_config.yaml'),
        nodes_list, postgresql_cluster, local_config_files)
    _prepare_rabbitmq_config_files(
        templates_env.get_template('rabbitmq_config.yaml'),
        nodes_list, rabbitmq_cluster, local_config_files)

def _prepare_manager_config_files(template, nodes_list, rabbitmq_cluster,
                                  postgresql_cluster, local_config_files):
    for i, node in enumerate(nodes_list, start=1):
        rendered_date = template.render(
            node=node,
            ca_path=join(REMOTE_CERTS_PATH, 'ca.pem'),
            license_path=REMOTE_LICENSE_PATH,
            rabbitmq_cluster=rabbitmq_cluster,
            postgresql_cluster=postgresql_cluster
        )
        config_name = 'manager-{0}_config.yaml'.format(i)
        remote_config_path = join(REMOTE_CONFIGS_PATH, config_name)
        local_config_file = local_config_files / config_name
        local_config_file.write_text(u'{0}'.format(rendered_date))

        nodes_list[0].put_remote_file(remote_config_path,
                                      str(local_config_file))
        node.local_manager_config_path = local_config_file
        node.remote_manager_config_path = remote_config_path

def _prepare_postgresql_config_files(template, nodes_list, postgresql_cluster,
                                     local_config_files):
    for i, node in enumerate(nodes_list, start=1):
        rendered_date = template.render(
            node=node,
            ca_path=join(REMOTE_CERTS_PATH, 'ca.pem'),
            postgresql_cluster=postgresql_cluster
        )
        config_name = 'postgresql-{0}_config.yaml'.format(i)
        remote_config_path = join(REMOTE_CONFIGS_PATH, config_name)
        local_config_file = local_config_files / config_name
        local_config_file.write_text(u'{0}'.format(rendered_date))

        nodes_list[0].put_remote_file(remote_config_path,
                                      str(local_config_file))
        node.local_postgresql_config_path = local_config_file
        node.remote_postgresql_config_path = remote_config_path

def _prepare_rabbitmq_config_files(template, nodes_list, rabbitmq_cluster,
                                   local_config_files):
    for i, node in enumerate(nodes_list, start=1):
        rendered_date = template.render(
            node=node,
            ca_path=join(REMOTE_CERTS_PATH, 'ca.pem'),
            rabbitmq_cluster=rabbitmq_cluster,
            join_cluster=nodes_list[0].hostname if i > 1 else None
        )
        config_name = 'rabbitmq-{0}_config.yaml'.format(i)
        remote_config_path = join(REMOTE_CONFIGS_PATH, config_name)
        local_config_file = local_config_files / config_name
        local_config_file.write_text(u'{0}'.format(rendered_date))

        nodes_list[0].put_remote_file(remote_config_path,
                                      str(local_config_file))
        node.local_rabbitmq_config_path = local_config_file
        node.remote_rabbitmq_config_path = remote_config_path

def _create_certificates(local_certs_path, nodes_list, pass_certs=False):
    ca_base = str(local_certs_path / 'ca.')
    ca_cert = ca_base + 'pem'
    ca_key = ca_base + 'key'
    generate_ca_cert(ca_cert, ca_key)
    for i, node in enumerate(nodes_list, start=1):
        node_cert = str(local_certs_path / 'node-{0}.crt'.format(i))
        node_key = str(local_certs_path / 'node-{0}.key'.format(i))
        generate_ssl_certificate(
            [node.private_ip_address, node.ip_address],
            node.hostname,
            node_cert,
            node_key,
            ca_cert,
            ca_key
        )
        if pass_certs:
            remote_cert = join(REMOTE_CERTS_PATH, 'node-{0}.crt'.format(i))
            remote_key = join(REMOTE_CERTS_PATH, 'node-{0}.key'.format(i))
            node.cert_path = remote_cert
            node.key_path = remote_key
            node.put_remote_file(remote_cert, node_cert)
            node.put_remote_file(remote_key, node_key)
            node.put_remote_file(join(REMOTE_CERTS_PATH, 'ca.pem'), ca_cert)

def _update_three_nodes_config_dict_vms(config_dict, existing_vms_list):
    for i, node in enumerate(existing_vms_list, start=1):
        config_dict['existing_vms']['node-{0}'.format(i)].update({
            'private_ip': str(node.private_ip_address),
            'public_ip': str(node.ip_address)
        })

def _install_cluster(node, config_dict, test_config, ssh_key):
    node.put_remote_file(remote_path=REMOTE_SSH_KEY_PATH,
                         local_path=ssh_key.private_key_path)

    node.put_remote_file(remote_path=REMOTE_LICENSE_PATH,
                         local_path=get_resource_path(
                             'test_valid_paying_license.yaml'))

    remote_cluster_config_path = '/tmp/cfy_cluster_config.yaml'
    node.put_remote_file_content(remote_cluster_config_path,
                                 yaml.dump(config_dict))

    node.run_command('yum install -y {0}'.format(
        test_config['cfy_cluster_manager']['rpm_path']), use_sudo=True)

    node.run_command('cfy_cluster_manager install -v --config-path '
                     '{0}'.format(remote_cluster_config_path))
<commit_msg>Adding using provided config files test<commit_after>import pkg_resources
from os.path import join

import yaml
import pytest
from jinja2 import Environment, FileSystemLoader

from cosmo_tester.framework.util import (generate_ca_cert,
                                         generate_ssl_certificate,
                                         get_resource_path)

RESOURCES_PATH = pkg_resources.resource_filename(
    'cosmo_tester', 'test_suites/cluster/cfy_cluster_manager_resources')
REMOTE_SSH_KEY_PATH = '/tmp/cfy_cluster_manager_ssh_key.pem'
REMOTE_LICENSE_PATH = '/tmp/cfy_cluster_manager_license.yaml'
REMOTE_CERTS_PATH = '/tmp/certs'
REMOTE_CONFIGS_PATH = '/tmp/config_files'

@pytest.fixture()
def basic_config_dict(ssh_key, test_config):
    return {
        'ssh_key_path': REMOTE_SSH_KEY_PATH,
        'ssh_user': 'centos',
        'cloudify_license_path': REMOTE_LICENSE_PATH,
        'manager_rpm_download_link': test_config['cfy_cluster_manager'][
            'manager_install_rpm_path']
    }

@pytest.fixture()
def three_nodes_config_dict(basic_config_dict):
    return _get_config_dict('three_nodes_config.yaml', basic_config_dict)

@pytest.fixture()
def nine_nodes_config_dict(basic_config_dict):
    return _get_config_dict('nine_nodes_config.yaml', basic_config_dict)

@pytest.fixture()
def local_certs_path(tmp_path):
    dir_path = tmp_path / 'certs'
    dir_path.mkdir()
    return dir_path

@pytest.fixture()
def local_config_files(tmp_path):
    dir_path = tmp_path / 'config_files'
    dir_path.mkdir()
    return dir_path

def _get_config_dict(config_file_name, basic_config_dict):
    config_path = join(RESOURCES_PATH, config_file_name)
    with open(config_path) as config_file:
        config_dict = yaml.safe_load(config_file)

    config_dict.update(basic_config_dict)
    return config_dict

def test_create_three_nodes_cluster(three_vms, three_nodes_config_dict,
                                    test_config, ssh_key):
    """Tests that a three nodes cluster is successfully created."""
    node1, node2, node3 = three_vms
    _update_three_nodes_config_dict_vms(three_nodes_config_dict,
                                        [node1, node2, node3])

    _install_cluster(node1, three_nodes_config_dict, test_config, ssh_key)

def test_create_nine_nodes_cluster(nine_vms, nine_nodes_config_dict,
                                   test_config, ssh_key):
    """Tests that a nine nodes cluster is successfully created."""
    node1, node2, node3, node4, node5, node6, node7, node8, node9 = nine_vms
    for i, node in enumerate([node1, node2, node3, node4, node5, node6,
                              node7, node8, node9]):
        node_num = (i % 3) + 1
        if i < 3:
            node_name = 'manager-{0}'.format(node_num)
        elif i < 6:
            node_name = 'rabbitmq-{0}'.format(node_num)
        else:
            node_name = 'postgresql-{0}'.format(node_num)

        nine_nodes_config_dict['existing_vms'][node_name].update({
            'private_ip': str(node.private_ip_address),
            'public_ip': str(node.ip_address)
        })

    _install_cluster(node1, nine_nodes_config_dict, test_config, ssh_key)

def test_create_three_nodes_cluster_using_certificates(
        three_vms, three_nodes_config_dict, test_config,
        ssh_key, local_certs_path, logger):
    """Tests that the supllied certificates are being used in the cluster."""
    node1, node2, node3 = three_vms
    nodes_list = [node1, node2, node3]

    logger.info('Creating certificates')
    _create_certificates(local_certs_path, nodes_list)

    logger.info('Copying certificates to node-1')
    node1.run_command('mkdir -p {0}'.format(REMOTE_CERTS_PATH))
    for cert in local_certs_path.iterdir():
        node1.put_remote_file(local_path=str(cert),
                              remote_path=join(REMOTE_CERTS_PATH, cert.name))

    logger.info('Preparing cluster install configuration file')
    _update_three_nodes_config_dict_vms(three_nodes_config_dict, nodes_list)
    three_nodes_config_dict['ca_cert_path'] = join(REMOTE_CERTS_PATH, 'ca.pem')
    for i, node in enumerate(nodes_list, start=1):
        three_nodes_config_dict['existing_vms']['node-{0}'.format(i)].update({
            'cert_path': join(REMOTE_CERTS_PATH, 'node-{0}.crt'.format(i)),
            'key_path': join(REMOTE_CERTS_PATH, 'node-{0}.key'.format(i))
        })

    logger.info('Installing cluster')
    _install_cluster(node1, three_nodes_config_dict, test_config, ssh_key)

    logger.info('Asserting certs were successfully copied')
    local_ca_path = local_certs_path / 'ca.pem'
    ca_path_in_use = '/etc/cloudify/ssl/cloudify_internal_ca_cert.pem'
    for i, node in enumerate(nodes_list, start=1):
        node_name = 'node-{0}'.format(i)
        logger.info('Asserting certificates for %s', node_name)
        local_node_cert_path = local_certs_path / '{0}.crt'.format(node_name)
        local_node_key_path = local_certs_path / '{0}.key'.format(node_name)
        cert_path_in_use = '/etc/cloudify/ssl/cloudify_internal_cert.pem'
        key_path_in_use = '/etc/cloudify/ssl/cloudify_internal_key.pem'

        assert (local_node_cert_path.read_text() ==
                node.get_remote_file_content(cert_path_in_use))

        assert (local_node_key_path.read_text() ==
                node.get_remote_file_content(key_path_in_use))

        assert (local_ca_path.read_text() ==
                node.get_remote_file_content(ca_path_in_use))

def test_three_nodes_using_provided_config_files(
        three_vms, three_nodes_config_dict, test_config,
        ssh_key, local_certs_path, local_config_files, logger):
    node1, node2, node3 = three_vms
    nodes_list = [node1, node2, node3]
    logger.info('Creating certificates and passing them to the instances')
    node1.run_command('mkdir -p {0}'.format(REMOTE_CERTS_PATH))
    _create_certificates(local_certs_path, nodes_list, pass_certs=True)

    logger.info('Preparing config files')
    _prepare_config_files(nodes_list, local_config_files)
    _update_three_nodes_config_dict_vms(three_nodes_config_dict, nodes_list)
    for i, node in enumerate(nodes_list, start=1):
        three_nodes_config_dict['existing_vms']['node-{0}'.format(i)][
            'config_path'].update({
                'manager_config_path': node.remote_manager_config_path,
                'postgresql_config_path': node.remote_postgresql_config_path,
                'rabbitmq_config_path': node.remote_rabbitmq_config_path
            })

    logger.info('Installing cluster')
    _install_cluster(node1, three_nodes_config_dict, test_config, ssh_key)

    logger.info('Asserting config_files')
    base_cfy_dir = '/etc/cloudify'
    for i, node in enumerate(nodes_list, start=1):
        logger.info(
            'Asserting config.yaml files for {0}'.format('node-{0}'.format(i)))

        assert (node.local_manager_config_path.read_text() ==
                node.get_remote_file_content(
                    join(base_cfy_dir, node.local_manager_config_path.name))
                )

        assert (node.local_postgresql_config_path.read_text() ==
                node.get_remote_file_content(
                    join(base_cfy_dir, node.local_postgresql_config_path.name))
                )

        assert (node.local_rabbitmq_config_path.read_text() ==
                node.get_remote_file_content(
                    join(base_cfy_dir, node.local_rabbitmq_config_path.name))
                )

def _prepare_config_files(nodes_list, local_config_files):
    rabbitmq_cluster = {
        node.hostname: {
            'networks': {
                'default': str(node.private_ip_address)
            }
        } for node in nodes_list
    }

    postgresql_cluster = {
        node.hostname: {
            'ip': str(node.private_ip_address)
        } for node in nodes_list
    }

    nodes_list[0].run_command('mkdir -p {0}'.format(REMOTE_CONFIGS_PATH))
    templates_env = Environment(loader=FileSystemLoader(
        join(RESOURCES_PATH, 'config_files_templates')))
    _prepare_manager_config_files(
        templates_env.get_template('manager_config.yaml'),
        nodes_list, rabbitmq_cluster, postgresql_cluster, local_config_files)
    _prepare_postgresql_config_files(
        templates_env.get_template('postgresql_config.yaml'),
        nodes_list, postgresql_cluster, local_config_files)
    _prepare_rabbitmq_config_files(
        templates_env.get_template('rabbitmq_config.yaml'),
        nodes_list, rabbitmq_cluster, local_config_files)

def _prepare_manager_config_files(template, nodes_list, rabbitmq_cluster,
                                  postgresql_cluster, local_config_files):
    for i, node in enumerate(nodes_list, start=1):
        rendered_date = template.render(
            node=node,
            ca_path=join(REMOTE_CERTS_PATH, 'ca.pem'),
            license_path=REMOTE_LICENSE_PATH,
            rabbitmq_cluster=rabbitmq_cluster,
            postgresql_cluster=postgresql_cluster
        )
        config_name = 'manager-{0}_config.yaml'.format(i)
        remote_config_path = join(REMOTE_CONFIGS_PATH, config_name)
        local_config_file = local_config_files / config_name
        local_config_file.write_text(u'{0}'.format(rendered_date))

        nodes_list[0].put_remote_file(remote_config_path,
                                      str(local_config_file))
        node.local_manager_config_path = local_config_file
        node.remote_manager_config_path = remote_config_path

def _prepare_postgresql_config_files(template, nodes_list, postgresql_cluster,
                                     local_config_files):
    for i, node in enumerate(nodes_list, start=1):
        rendered_date = template.render(
            node=node,
            ca_path=join(REMOTE_CERTS_PATH, 'ca.pem'),
            postgresql_cluster=postgresql_cluster
        )
        config_name = 'postgresql-{0}_config.yaml'.format(i)
        remote_config_path = join(REMOTE_CONFIGS_PATH, config_name)
        local_config_file = local_config_files / config_name
        local_config_file.write_text(u'{0}'.format(rendered_date))

        nodes_list[0].put_remote_file(remote_config_path,
                                      str(local_config_file))
        node.local_postgresql_config_path = local_config_file
        node.remote_postgresql_config_path = remote_config_path

def _prepare_rabbitmq_config_files(template, nodes_list, rabbitmq_cluster,
                                   local_config_files):
    for i, node in enumerate(nodes_list, start=1):
        rendered_date = template.render(
            node=node,
            ca_path=join(REMOTE_CERTS_PATH, 'ca.pem'),
            rabbitmq_cluster=rabbitmq_cluster,
            join_cluster=nodes_list[0].hostname if i > 1 else None
        )
        config_name = 'rabbitmq-{0}_config.yaml'.format(i)
        remote_config_path = join(REMOTE_CONFIGS_PATH, config_name)
        local_config_file = local_config_files / config_name
        local_config_file.write_text(u'{0}'.format(rendered_date))

        nodes_list[0].put_remote_file(remote_config_path,
                                      str(local_config_file))
        node.local_rabbitmq_config_path = local_config_file
        node.remote_rabbitmq_config_path = remote_config_path

def _create_certificates(local_certs_path, nodes_list, pass_certs=False):
    ca_base = str(local_certs_path / 'ca.')
    ca_cert = ca_base + 'pem'
    ca_key = ca_base + 'key'
    generate_ca_cert(ca_cert, ca_key)
    for i, node in enumerate(nodes_list, start=1):
        node_cert = str(local_certs_path / 'node-{0}.crt'.format(i))
        node_key = str(local_certs_path / 'node-{0}.key'.format(i))
        generate_ssl_certificate(
            [node.private_ip_address, node.ip_address],
            node.hostname,
            node_cert,
            node_key,
            ca_cert,
            ca_key
        )
        if pass_certs:
            remote_cert = join(REMOTE_CERTS_PATH, 'node-{0}.crt'.format(i))
            remote_key = join(REMOTE_CERTS_PATH, 'node-{0}.key'.format(i))
            node.cert_path = remote_cert
            node.key_path = remote_key
            node.put_remote_file(remote_cert, node_cert)
            node.put_remote_file(remote_key, node_key)
            node.put_remote_file(join(REMOTE_CERTS_PATH, 'ca.pem'), ca_cert)

def _update_three_nodes_config_dict_vms(config_dict, existing_vms_list):
    for i, node in enumerate(existing_vms_list, start=1):
        config_dict['existing_vms']['node-{0}'.format(i)].update({
            'private_ip': str(node.private_ip_address),
            'public_ip': str(node.ip_address)
        })

def _install_cluster(node, config_dict, test_config, ssh_key):
    node.put_remote_file(remote_path=REMOTE_SSH_KEY_PATH,
                         local_path=ssh_key.private_key_path)

    node.put_remote_file(remote_path=REMOTE_LICENSE_PATH,
                         local_path=get_resource_path(
                             'test_valid_paying_license.yaml'))

    remote_cluster_config_path = '/tmp/cfy_cluster_config.yaml'
    node.put_remote_file_content(remote_cluster_config_path,
                                 yaml.dump(config_dict))

    node.run_command('yum install -y {0}'.format(
        test_config['cfy_cluster_manager']['rpm_path']), use_sudo=True)

    node.run_command('cfy_cluster_manager install -v --config-path '
                     '{0}'.format(remote_cluster_config_path))
<|endoftext|>


<commit_before>from multiprocessing import Pool
import h5py as hdf
import numpy as np
from data_handler import mkTruth, mkHalo, mkQs
from halo_handler import find_indices, find_indices_multi
from calc_cluster_props import (updateArray, findClusterRedshift, findLOSV,
                                findLOSVDmcmc, calc_mass_Evrard)
import os
import random

class AsyncFactory:
    def __init__(self, func, cb_func):
        self.func = func
        self.cb_func = cb_func
        self.pool = Pool(maxtasksperchild=10)

    def call(self,*args, **kwargs):
        self.pool.apply_async(self.func, args, kwargs, self.cb_func)

    def wait(self):
        self.pool.close()
        self.pool.join()

def randomR():
    ''' randomly choose a location on the unit hemisphere, r = [x,y,z] with
    -1<x<1, -1<y<1, and 0<z<1 and x^2+y^2+z^2 = 1.

    '''

    phi = 2*np.pi*random.random()
    theta = np.arccos(np.random.random())
    r = [np.sin(theta)*np.cos(phi), np.sin(theta)*np.sin(phi), np.cos(theta)]
    return r

def worker(pos, data, center, rot=False):
    #print "PID: %d \t Value: %d" % (os.getpid(), pos)
    data = updateArray(data)
    data = findClusterRedshift(data)
    #data = findSeperationSpatial(data, center)

    # do the rotations or not
    if rot:
        r = randomR()
        v = np.column_stack([data['VX'], data['VY'], data['VZ']])
        rot = [np.dot(r,vi) for vi in v]
        data['LOSV'] = rot
    else:
        data = findLOSV(data)

    data, sigma_dist = findLOSVDmcmc(data)
    data = calc_mass_Evrard(data, A1D = 1177, alpha = 0.364)
    return pos, data, sigma_dist

def cb_func((pos, data, sigma_dist)):
    if pos % 200 == 0:
        print pos
    results['IDX'][pos] = pos
    results['CLUSZ'][pos] = data['CLUSZ'][0]
    results['LOSVD'][pos] = data['LOSVD'][0]
    results['MASS'][pos] = data['MASS'][0]
    results['LOSVD_err'][pos] = data['LOSVD_err'][0]
    results['LOSVD_dist'][pos] = sigma_dist[:,0]

if __name__ == "__main__":
    async_worker = AsyncFactory(worker, cb_func)
    halo = mkHalo()
    truth = mkTruth()
    Qs = mkQs()

    # this is the part that makes it realistic or not
    mask = (Qs == 0) | (Qs == 1)
    truth = truth[mask]

    # there are no clusters with mass < 2e11 and more than 5 galaxies
    mask = (halo['upid'] == -1) & (halo['m200c'] > 2e11)
    maskedHalo = halo[mask]
    hids, uniqueIdx = np.unique(maskedHalo['id'], return_index=True)

    # now we find all member halos of the most massive halos
    subHalos = find_indices(halo['upid'], hids)

    # now find all the corresponding galaxies
    gals = find_indices_multi(truth['HALOID'], halo['id'], subHalos)

    # figure out which clusters have an appropriate size
    x = np.array([i for i,g in enumerate(gals) if g.size >=5])

    #### Here is where we are going to make a flat HMF ####
    massBins = np.arange(11.5, 16, 0.5)
    masses = [np.log10(maskedHalo['m200c'][uniqueIdx[i]]/0.70) for i in x]
    indexes = np.digitize(masses, massBins)
    # now we choose some of the indexes uniformly -- How many do we want?
    indxChoice = np.random.randint(1, massBins.size, 1000)

    # figure out where they fit, and pick one, also uniformly
    picks = [np.random.choice(np.where(indexes==i)[0]) for i in indxChoice]

    # make the results container
    results = np.zeros((x[picks].size,), dtype=[('IDX', '>i4'),
        ('HALOID', '>i8'),
        ('ZSPEC', '>f4'),
        ('M200c', '>f4'),
        ('CLUSZ', '>f4'),
        ('LOSVD', '>f4'),
        ('MASS', '>f4'),
        ('NGAL', '>i4'),
        ('LOSVD_err', '>f4', (2,)),
        ('LOSVD_dist', '>f4', (10000,))])

    print('do work', x[picks].size, 'clusters to go!')
    for j,i in enumerate(x[picks]):
        center = (maskedHalo['ra'][uniqueIdx[i]],
                maskedHalo['dec'][uniqueIdx[i]])

        # start the rotations
        #if j > len(x):
        #    async_worker.call(j, truth[gals[i]], center, rot=True)
        #else:
        async_worker.call(j, truth[gals[i]], center, rot=True)
        #worker(j, truth[gals[i]], center, rot=False)

        # update results array
        results['HALOID'][j] = maskedHalo['id'][uniqueIdx[i]]
        results['NGAL'][j] = gals[i].size
        results['ZSPEC'][j] = maskedHalo['zspec'][uniqueIdx[i]]
        # millenium cosmology
        #results['M200c'][j] = maskedHalo['m200c'][uniqueIdx[i]]/0.73
        results['M200c'][j] = maskedHalo['m200c'][uniqueIdx[i]]/0.70

    async_worker.wait()

    print('results')
    try:
        os.remove('surveyComplete_noRotations.hdf5')
    except OSError:
        pass
    with hdf.File('surveyComplete_noRotations.hdf5', 'w') as f:
        f['surveyComplete'] = results
<commit_msg>fixed typo and updated number<commit_after>from multiprocessing import Pool
import h5py as hdf
import numpy as np
from data_handler import mkTruth, mkHalo, mkQs
from halo_handler import find_indices, find_indices_multi
from calc_cluster_props import (updateArray, findClusterRedshift, findLOSV,
                                findLOSVDmcmc, calc_mass_Evrard)
import os
import random

class AsyncFactory:
    def __init__(self, func, cb_func):
        self.func = func
        self.cb_func = cb_func
        self.pool = Pool(maxtasksperchild=10)

    def call(self,*args, **kwargs):
        self.pool.apply_async(self.func, args, kwargs, self.cb_func)

    def wait(self):
        self.pool.close()
        self.pool.join()

def randomR():
    ''' randomly choose a location on the unit hemisphere, r = [x,y,z] with
    -1<x<1, -1<y<1, and 0<z<1 and x^2+y^2+z^2 = 1.

    '''

    phi = 2*np.pi*random.random()
    theta = np.arccos(np.random.random())
    r = [np.sin(theta)*np.cos(phi), np.sin(theta)*np.sin(phi), np.cos(theta)]
    return r

def worker(pos, data, center, rot=False):
    #print "PID: %d \t Value: %d" % (os.getpid(), pos)
    data = updateArray(data)
    data = findClusterRedshift(data)
    #data = findSeperationSpatial(data, center)

    # do the rotations or not
    if rot:
        r = randomR()
        v = np.column_stack([data['VX'], data['VY'], data['VZ']])
        rot = [np.dot(r,vi) for vi in v]
        data['LOSV'] = rot
    else:
        data = findLOSV(data)

    data, sigma_dist = findLOSVDmcmc(data)
    data = calc_mass_Evrard(data, A1D = 1177, alpha = 0.364)
    return pos, data, sigma_dist

def cb_func((pos, data, sigma_dist)):
    if pos % 200 == 0:
        print pos
    results['IDX'][pos] = pos
    results['CLUSZ'][pos] = data['CLUSZ'][0]
    results['LOSVD'][pos] = data['LOSVD'][0]
    results['MASS'][pos] = data['MASS'][0]
    results['LOSVD_err'][pos] = data['LOSVD_err'][0]
    results['LOSVD_dist'][pos] = sigma_dist[:,0]

if __name__ == "__main__":
    async_worker = AsyncFactory(worker, cb_func)
    halo = mkHalo()
    truth = mkTruth(flatHMF=True)
    Qs = mkQs()

    # this is the part that makes it realistic or not
    mask = (Qs == 0) | (Qs == 1)
    truth = truth[mask]

    # there are no clusters with mass < 2e11 and more than 5 galaxies
    mask = (halo['upid'] == -1) & (halo['m200c'] > 2e11)
    maskedHalo = halo[mask]
    hids, uniqueIdx = np.unique(maskedHalo['id'], return_index=True)

    # now we find all member halos of the most massive halos
    subHalos = find_indices(halo['upid'], hids)

    # now find all the corresponding galaxies
    gals = find_indices_multi(truth['HALOID'], halo['id'], subHalos)

    # figure out which clusters have an appropriate size
    x = np.array([i for i,g in enumerate(gals) if g.size >=5])

    #### Here is where we are going to make a flat HMF ####
    massBins = np.arange(11.5, 16, 0.5)
    masses = [np.log10(maskedHalo['m200c'][uniqueIdx[i]]/0.70) for i in x]
    indexes = np.digitize(masses, massBins)
    # now we choose some of the indexes uniformly -- How many do we want?
    indxChoice = np.random.randint(1, massBins.size, 5000)

    # figure out where they fit, and pick one, also uniformly
    picks = [np.random.choice(np.where(indexes==i)[0]) for i in indxChoice]

    # make the results container
    results = np.zeros((x[picks].size,), dtype=[('IDX', '>i4'),
        ('HALOID', '>i8'),
        ('ZSPEC', '>f4'),
        ('M200c', '>f4'),
        ('CLUSZ', '>f4'),
        ('LOSVD', '>f4'),
        ('MASS', '>f4'),
        ('NGAL', '>i4'),
        ('LOSVD_err', '>f4', (2,)),
        ('LOSVD_dist', '>f4', (10000,))])

    print('do work', x[picks].size, 'clusters to go!')
    for j,i in enumerate(x[picks]):
        center = (maskedHalo['ra'][uniqueIdx[i]],
                maskedHalo['dec'][uniqueIdx[i]])

        # start the rotations
        #if j > len(x):
        #    async_worker.call(j, truth[gals[i]], center, rot=True)
        #else:
        async_worker.call(j, truth[gals[i]], center, rot=True)
        #worker(j, truth[gals[i]], center, rot=False)

        # update results array
        results['HALOID'][j] = maskedHalo['id'][uniqueIdx[i]]
        results['NGAL'][j] = gals[i].size
        results['ZSPEC'][j] = maskedHalo['zspec'][uniqueIdx[i]]
        # millenium cosmology
        #results['M200c'][j] = maskedHalo['m200c'][uniqueIdx[i]]/0.73
        results['M200c'][j] = maskedHalo['m200c'][uniqueIdx[i]]/0.70

    async_worker.wait()

    print('results')
    with hdf.File('buzzard_targetedRealistic_flatHMF.hdf5', 'w') as f:
        f['surveyComplete'] = results
<|endoftext|>


<commit_before># (C) Datadog, Inc. 2018
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
import os
import re
from datetime import datetime
from uuid import uuid4

from ..utils import (
    create_file,
    dir_exists,
    ensure_parent_dir_exists,
    path_join,
    read_file,
    read_file_binary,
    write_file,
    write_file_binary,
)
from .utils import normalize_package_name

TEMPLATES_DIR = path_join(os.path.dirname(os.path.abspath(__file__)), 'templates')
BINARY_EXTENSIONS = ('.png',)
SIMPLE_NAME = r'^\w+$'

def get_valid_templates():
    return sorted(os.listdir(TEMPLATES_DIR))

def construct_template_fields(integration_name, repo_choice, **kwargs):
    normalized_integration_name = normalize_package_name(integration_name)
    check_name_kebab = re.sub('_| ', '-', integration_name.lower())

    if repo_choice == 'core':
        author = 'Datadog'
        email = 'help@datadoghq.com'
        email_packages = 'packages@datadoghq.com'
        install_info = (
            'The {integration_name} check is included in the [Datadog Agent][2] package.\n'
            'No additional installation is needed on your server.'.format(integration_name=integration_name)
        )
        license_header = (
            '# (C) Datadog, Inc. {year}\n'
            '# All rights reserved\n'
            '# Licensed under a 3-clause BSD style license (see LICENSE)\n'.format(year=str(datetime.now().year))
        )
        support_type = 'core'
        test_dev_dep = '-e ../datadog_checks_dev'
        tox_base_dep = '-e../datadog_checks_base[deps]'
    else:
        author = 'U.N. Owen'
        email = email_packages = 'friend@datadog.community'
        install_info = (
            'The {} check is not included in the [Datadog Agent][2] package, so it must\n'
            'be installed manually.'.format(integration_name)
        )
        license_header = ''
        support_type = 'contrib'
        test_dev_dep = 'datadog-checks-dev'
        tox_base_dep = 'datadog-checks-base[deps]>=6.6.0'

    config = {
        'author': author,
        'check_class': '{}Check'.format(''.join(part.capitalize() for part in normalized_integration_name.split('_'))),
        'check_name': normalized_integration_name,
        'integration_name': integration_name,
        'check_name_kebab': check_name_kebab,
        'email': email,
        'email_packages': email_packages,
        'guid': uuid4(),
        'license_header': license_header,
        'install_info': install_info,
        'repo_choice': repo_choice,
        'support_type': support_type,
        'test_dev_dep': test_dev_dep,
        'tox_base_dep': tox_base_dep,
    }
    config.update(kwargs)

    return config

def create_template_files(template_name, new_root, config, read=False):
    files = []

    template_root = path_join(TEMPLATES_DIR, template_name)
    if not dir_exists(template_root):
        return files

    for root, _, template_files in os.walk(template_root):
        for template_file in template_files:
            if not template_file.endswith(('.pyc', '.pyo')):
                template_path = path_join(root, template_file)

                file_path = template_path.replace(template_root, '')
                file_path = '{}{}'.format(new_root, file_path.format(**config))

                files.append(File(file_path, template_path, config, read=read))

    return files

class File(object):
    def __init__(self, file_path, template_path, config, read=False):
        self.file_path = file_path
        self.template_path = template_path
        self.config = config
        self.binary = template_path.endswith(BINARY_EXTENSIONS)
        self._read = read_file_binary if self.binary else read_file
        self._write = write_file_binary if self.binary else write_file
        self.contents = None

        if read:
            self.read()

    def read(self):
        contents = self._read(self.template_path)

        if self.binary:
            self.contents = contents
        else:
            print(self.file_path)
            self.contents = contents.format(**self.config)

    def write(self):
        if self.contents is None:
            create_file(self.file_path)
        else:
            ensure_parent_dir_exists(self.file_path)
            self._write(self.file_path, self.contents)
<commit_msg>Remove spurious debug line (#3703)<commit_after># (C) Datadog, Inc. 2018
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
import os
import re
from datetime import datetime
from uuid import uuid4

from ..utils import (
    create_file,
    dir_exists,
    ensure_parent_dir_exists,
    path_join,
    read_file,
    read_file_binary,
    write_file,
    write_file_binary,
)
from .utils import normalize_package_name

TEMPLATES_DIR = path_join(os.path.dirname(os.path.abspath(__file__)), 'templates')
BINARY_EXTENSIONS = ('.png',)
SIMPLE_NAME = r'^\w+$'

def get_valid_templates():
    return sorted(os.listdir(TEMPLATES_DIR))

def construct_template_fields(integration_name, repo_choice, **kwargs):
    normalized_integration_name = normalize_package_name(integration_name)
    check_name_kebab = re.sub('_| ', '-', integration_name.lower())

    if repo_choice == 'core':
        author = 'Datadog'
        email = 'help@datadoghq.com'
        email_packages = 'packages@datadoghq.com'
        install_info = (
            'The {integration_name} check is included in the [Datadog Agent][2] package.\n'
            'No additional installation is needed on your server.'.format(integration_name=integration_name)
        )
        license_header = (
            '# (C) Datadog, Inc. {year}\n'
            '# All rights reserved\n'
            '# Licensed under a 3-clause BSD style license (see LICENSE)\n'.format(year=str(datetime.now().year))
        )
        support_type = 'core'
        test_dev_dep = '-e ../datadog_checks_dev'
        tox_base_dep = '-e../datadog_checks_base[deps]'
    else:
        author = 'U.N. Owen'
        email = email_packages = 'friend@datadog.community'
        install_info = (
            'The {} check is not included in the [Datadog Agent][2] package, so it must\n'
            'be installed manually.'.format(integration_name)
        )
        license_header = ''
        support_type = 'contrib'
        test_dev_dep = 'datadog-checks-dev'
        tox_base_dep = 'datadog-checks-base[deps]>=6.6.0'

    config = {
        'author': author,
        'check_class': '{}Check'.format(''.join(part.capitalize() for part in normalized_integration_name.split('_'))),
        'check_name': normalized_integration_name,
        'integration_name': integration_name,
        'check_name_kebab': check_name_kebab,
        'email': email,
        'email_packages': email_packages,
        'guid': uuid4(),
        'license_header': license_header,
        'install_info': install_info,
        'repo_choice': repo_choice,
        'support_type': support_type,
        'test_dev_dep': test_dev_dep,
        'tox_base_dep': tox_base_dep,
    }
    config.update(kwargs)

    return config

def create_template_files(template_name, new_root, config, read=False):
    files = []

    template_root = path_join(TEMPLATES_DIR, template_name)
    if not dir_exists(template_root):
        return files

    for root, _, template_files in os.walk(template_root):
        for template_file in template_files:
            if not template_file.endswith(('.pyc', '.pyo')):
                template_path = path_join(root, template_file)

                file_path = template_path.replace(template_root, '')
                file_path = '{}{}'.format(new_root, file_path.format(**config))

                files.append(File(file_path, template_path, config, read=read))

    return files

class File(object):
    def __init__(self, file_path, template_path, config, read=False):
        self.file_path = file_path
        self.template_path = template_path
        self.config = config
        self.binary = template_path.endswith(BINARY_EXTENSIONS)
        self._read = read_file_binary if self.binary else read_file
        self._write = write_file_binary if self.binary else write_file
        self.contents = None

        if read:
            self.read()

    def read(self):
        contents = self._read(self.template_path)

        if self.binary:
            self.contents = contents
        else:
            self.contents = contents.format(**self.config)

    def write(self):
        if self.contents is None:
            create_file(self.file_path)
        else:
            ensure_parent_dir_exists(self.file_path)
            self._write(self.file_path, self.contents)
<|endoftext|>


<commit_before>""" Create a Twitter SocialUser object from userid or screen_name """

import logging
from django.core.management.base import BaseCommand, CommandError
from moderation.profile import create_twitter_social_user_and_profile
from community.models import Community
from moderation.models import SocialUser, Category, UserCategoryRelationship
from bot.tweepy_api import get_api
from tweepy.error import TweepError
from typing import List
from moderation.social import get_socialuser_from_screen_name

logger = logging.getLogger(__name__)

class Command(BaseCommand):
    help = 'Create SocialUser object from userids or screen names'

    def add_arguments(self, parser):
        parser.add_argument('users', nargs='+', type=str)
        parser.add_argument(
            "-c",
            "--category",
            type=str,
            help="category name"
        )
        parser.add_argument(
            "-m",
            "--moderator",
            type=str,
            help="moderator screen_name"
        )
        parser.add_argument(
            "-k",
            "--community",
            type=str,
            help="community name"
        )

    def create_user_category_relationship(self, su):
        try:
            UserCategoryRelationship.objects.create(
                social_user=su,
                category=self.category,
                moderator=self.moderator,
                community=self.community,
            )
        except:
            return

    def process_user_id(self, user_id):
        try:
            su = SocialUser.objects.get(user_id=user_id)
            self.stdout.write(
                self.style.WARNING(
                    f'SocialUser object with user id {user_id} already exists.'
                )
            )
            if self.categorize:
                self.create_user_category_relationship(su)
        except SocialUser.DoesNotExist:
            return

    def process_screen_name(self, screen_name: str):
        api = get_api()
        if not api:
            self.stdout.write(
                self.style.ERROR(
                    f'Could not get functional API.'
                )
            )
            return
        try:
            tweepy_user = get_api().get_user(screen_name=screen_name)
        except TweepError as e:
            raise CommandError('Tweepy error "%s' % e)
        except AttributeError as e:
            logger.error(f"AttributeError: {e}")
            return
        user_id=tweepy_user.id
        self.stdout.write(f"user id: {user_id}")
        su, created = create_twitter_social_user_and_profile(user_id)
        if created and su:
            su.refresh_from_db()
            self.stdout.write(
                self.style.SUCCESS(
                    f'Done creating SocialUser object {su}.'
                )
            )
        elif su:
            self.stdout.write(
                self.style.SUCCESS(
                    f'SocialUser object {su} already exists. '
                    'Its profile was updated.'
                )
            )
            current_screen_name = su.screen_name_tag()
            if not (current_screen_name == screen_name):
                self.stdout.write(
                    self.style.WARNING(
                        f'screen_name of SocialUser {su} has changed!\n'
                        f'old: "{screen_name}" -> new: "{current_screen_name}"'
                    )
                )
        else:
            self.stdout.write(
                self.style.ERROR(
                    f'Error, no SocialUser was created.'
                )
            )
            return
        if su and self.categorize:
            self.create_user_category_relationship(su)

    def handle(self, *args, **options):
        users: List =  options['users']
        mod_screen_name: str =  options['moderator']
        category: str =  options['category']
        community: str =  options['community']
        try:
            self.category=Category.objects.get(name=category)
            logger.debug(self.category)
        except Category.DoesNotExist:
            self.category=None
        try:
            self.community=Community.objects.get(name=community)
            logger.debug(self.community)
        except Community.DoesNotExist:
            self.community=None
        if mod_screen_name:
            self.moderator = get_socialuser_from_screen_name(mod_screen_name)
            if not self.moderator:
                self.stdout.write(
                    self.style.ERROR(
                        f'SocialUser {mod_screen_name=} does not exist.'
                    )
                )
                return
        self.categorize: bool = (
            self.category
            and self.moderator
            and self.community
        )
        if not users:
            self.stdout.write(
                self.style.ERROR(
                    'You must provide one or more username(s) or userid(s).'
                )
            )
            return
        for user in users:
            try:
                user_id=int(user)
                self.process_user_id(user_id)
            except ValueError:
                self.process_screen_name(user)<commit_msg>fix command (old, new screen_name)<commit_after>""" Create a Twitter SocialUser object from userid or screen_name """

import logging
from django.core.management.base import BaseCommand, CommandError
from moderation.profile import create_twitter_social_user_and_profile
from community.models import Community
from moderation.models import SocialUser, Category, UserCategoryRelationship
from bot.tweepy_api import get_api
from tweepy.error import TweepError
from typing import List
from moderation.social import get_socialuser_from_screen_name

logger = logging.getLogger(__name__)

class Command(BaseCommand):
    help = 'Create SocialUser object from userids or screen names'

    def add_arguments(self, parser):
        parser.add_argument('users', nargs='+', type=str)
        parser.add_argument(
            "-c",
            "--category",
            type=str,
            help="category name"
        )
        parser.add_argument(
            "-m",
            "--moderator",
            type=str,
            help="moderator screen_name"
        )
        parser.add_argument(
            "-k",
            "--community",
            type=str,
            help="community name"
        )

    def create_user_category_relationship(self, su):
        try:
            UserCategoryRelationship.objects.create(
                social_user=su,
                category=self.category,
                moderator=self.moderator,
                community=self.community,
            )
        except:
            return

    def process_user_id(self, user_id):
        try:
            su = SocialUser.objects.get(user_id=user_id)
            self.stdout.write(
                self.style.WARNING(
                    f'SocialUser object with user id {user_id} already exists.'
                )
            )
            if self.categorize:
                self.create_user_category_relationship(su)
        except SocialUser.DoesNotExist:
            return

    def process_screen_name(self, screen_name: str):
        api = get_api()
        if not api:
            self.stdout.write(
                self.style.ERROR(
                    f'Could not get functional API.'
                )
            )
            return
        try:
            tweepy_user = get_api().get_user(screen_name=screen_name)
        except TweepError as e:
            raise CommandError('Tweepy error "%s' % e)
        except AttributeError as e:
            logger.error(f"AttributeError: {e}")
            return
        user_id=tweepy_user.id
        self.stdout.write(f"user id: {user_id}")
        su, created = create_twitter_social_user_and_profile(user_id)
        if created and su:
            su.refresh_from_db()
            self.stdout.write(
                self.style.SUCCESS(
                    f'Done creating SocialUser object {su}.'
                )
            )
        elif su:
            self.stdout.write(
                self.style.SUCCESS(
                    f'SocialUser object {su} already exists. '
                    'Its profile was updated.'
                )
            )
            old_s_n = su.screen_name_tag()
            su.refresh_from_db()
            current_s_n = su.screen_name_tag()
            if not (current_s_n == old_s_n):
                self.stdout.write(
                    self.style.WARNING(
                        f'screen_name of SocialUser {su} has changed!\n'
                        f'old: "{old_s_n}" -> new: "{current_s_n}"'
                    )
                )
        else:
            self.stdout.write(
                self.style.ERROR(
                    f'Error, no SocialUser was created.'
                )
            )
            return
        if su and self.categorize:
            self.create_user_category_relationship(su)

    def handle(self, *args, **options):
        users: List =  options['users']
        mod_screen_name: str =  options['moderator']
        category: str =  options['category']
        community: str =  options['community']
        try:
            self.category=Category.objects.get(name=category)
            logger.debug(self.category)
        except Category.DoesNotExist:
            self.category=None
        try:
            self.community=Community.objects.get(name=community)
            logger.debug(self.community)
        except Community.DoesNotExist:
            self.community=None
        if mod_screen_name:
            self.moderator = get_socialuser_from_screen_name(mod_screen_name)
            if not self.moderator:
                self.stdout.write(
                    self.style.ERROR(
                        f'SocialUser {mod_screen_name=} does not exist.'
                    )
                )
                return
        self.categorize: bool = (
            self.category
            and self.moderator
            and self.community
        )
        if not users:
            self.stdout.write(
                self.style.ERROR(
                    'You must provide one or more username(s) or userid(s).'
                )
            )
            return
        for user in users:
            try:
                user_id=int(user)
                self.process_user_id(user_id)
            except ValueError:
                self.process_screen_name(user)<|endoftext|>


<commit_before>import asyncio
import json
import re

import pytest

pytest.importorskip("bokeh")

from tornado.escape import url_escape
from tornado.httpclient import AsyncHTTPClient, HTTPClientError

from dask.sizeof import sizeof

from distributed.utils import is_valid_xml
from distributed.utils_test import gen_cluster, inc, slowinc

@gen_cluster(client=True)
async def test_connect(c, s, a, b):
    future = c.submit(lambda x: x + 1, 1)
    x = c.submit(slowinc, 1, delay=1, retries=5)
    await future
    http_client = AsyncHTTPClient()
    for suffix in [
        "info/main/workers.html",
        "info/worker/" + url_escape(a.address) + ".html",
        "info/task/" + url_escape(future.key) + ".html",
        "info/main/logs.html",
        "info/logs/" + url_escape(a.address) + ".html",
        "info/call-stack/" + url_escape(x.key) + ".html",
        "info/call-stacks/" + url_escape(a.address) + ".html",
        "json/counts.json",
        "json/identity.json",
        "json/index.html",
        "individual-plots.json",
        "sitemap.json",
    ]:
        response = await http_client.fetch(
            "http://localhost:%d/%s" % (s.http_server.port, suffix)
        )
        assert response.code == 200
        body = response.body.decode()
        if suffix.endswith(".json"):
            json.loads(body)
        else:
            assert is_valid_xml(body)
            assert not re.search("href=./", body)  # no absolute links

@gen_cluster(client=True, nthreads=[])
async def test_worker_404(c, s):
    http_client = AsyncHTTPClient()
    with pytest.raises(HTTPClientError) as err:
        await http_client.fetch(
            "http://localhost:%d/info/worker/unknown" % s.http_server.port
        )
    assert err.value.code == 404
    with pytest.raises(HTTPClientError) as err:
        await http_client.fetch(
            "http://localhost:%d/info/task/unknown" % s.http_server.port
        )
    assert err.value.code == 404

@gen_cluster(client=True, scheduler_kwargs={"http_prefix": "/foo", "dashboard": True})
async def test_prefix(c, s, a, b):
    http_client = AsyncHTTPClient()
    for suffix in ["foo/info/main/workers.html", "foo/json/index.html", "foo/system"]:
        response = await http_client.fetch(
            "http://localhost:%d/%s" % (s.http_server.port, suffix)
        )
        assert response.code == 200
        body = response.body.decode()
        if suffix.endswith(".json"):
            json.loads(body)
        else:
            assert is_valid_xml(body)

@gen_cluster(client=True, clean_kwargs={"threads": False})
async def test_prometheus(c, s, a, b):
    pytest.importorskip("prometheus_client")
    from prometheus_client.parser import text_string_to_metric_families

    http_client = AsyncHTTPClient()

    # request data twice since there once was a case where metrics got registered multiple times resulting in
    # prometheus_client errors
    for _ in range(2):
        response = await http_client.fetch(
            "http://localhost:%d/metrics" % s.http_server.port
        )
        assert response.code == 200
        assert response.headers["Content-Type"] == "text/plain; version=0.0.4"

        txt = response.body.decode("utf8")
        families = {
            family.name: family for family in text_string_to_metric_families(txt)
        }
        assert "dask_scheduler_workers" in families

        client = families["dask_scheduler_clients"]
        assert client.samples[0].value == 1.0

@pytest.mark.repeat(100)
@gen_cluster(client=True, clean_kwargs={"threads": False})
async def test_prometheus_collect_task_states(c, s, a, b):
    pytest.importorskip("prometheus_client")
    from prometheus_client.parser import text_string_to_metric_families

    http_client = AsyncHTTPClient()

    async def fetch_metrics():
        port = s.http_server.port
        response = await http_client.fetch(f"http://localhost:{port}/metrics")
        txt = response.body.decode("utf8")
        families = {
            family.name: family for family in text_string_to_metric_families(txt)
        }

        active_metrics = {
            sample.labels["state"]: sample.value
            for sample in families["dask_scheduler_tasks"].samples
        }
        forgotten_tasks = [
            sample.value
            for sample in families["dask_scheduler_tasks_forgotten"].samples
        ]
        return active_metrics, forgotten_tasks

    expected = {"memory", "released", "processing", "waiting", "no-worker", "erred"}

    # Ensure that we get full zero metrics for all states even though the
    # scheduler did nothing, yet
    assert not s.tasks
    active_metrics, forgotten_tasks = await fetch_metrics()
    assert active_metrics.keys() == expected
    assert sum(active_metrics.values()) == 0.0
    assert sum(forgotten_tasks) == 0.0

    # submit a task which should show up in the prometheus scraping
    future = c.submit(slowinc, 1, delay=0.5)
    while not any(future.key in w.tasks for w in [a, b]):
        await asyncio.sleep(0.001)

    active_metrics, forgotten_tasks = await fetch_metrics()
    assert active_metrics.keys() == expected
    assert sum(active_metrics.values()) == 1.0
    assert sum(forgotten_tasks) == 0.0

    res = await c.gather(future)
    assert res == 2

    future.release()

    while any(future.key in w.tasks for w in [a, b]):
        await asyncio.sleep(0.001)

    active_metrics, forgotten_tasks = await fetch_metrics()
    assert active_metrics.keys() == expected
    assert sum(active_metrics.values()) == 0.0
    assert sum(forgotten_tasks) == 0.0

@gen_cluster(client=True, clean_kwargs={"threads": False})
async def test_health(c, s, a, b):
    http_client = AsyncHTTPClient()

    response = await http_client.fetch(
        "http://localhost:%d/health" % s.http_server.port
    )
    assert response.code == 200
    assert response.headers["Content-Type"] == "text/plain"

    txt = response.body.decode("utf8")
    assert txt == "ok"

@gen_cluster()
async def test_sitemap(s, a, b):
    http_client = AsyncHTTPClient()

    response = await http_client.fetch(
        "http://localhost:%d/sitemap.json" % s.http_server.port
    )
    out = json.loads(response.body.decode())
    assert "paths" in out
    assert "/sitemap.json" in out["paths"]
    assert "/health" in out["paths"]
    assert "/statics/css/base.css" in out["paths"]

@gen_cluster(client=True)
async def test_task_page(c, s, a, b):
    future = c.submit(lambda x: x + 1, 1, workers=a.address)
    x = c.submit(inc, 1)
    await future
    http_client = AsyncHTTPClient()

    "info/task/" + url_escape(future.key) + ".html",
    response = await http_client.fetch(
        "http://localhost:%d/info/task/" % s.http_server.port
        + url_escape(future.key)
        + ".html"
    )
    assert response.code == 200
    body = response.body.decode()

    assert str(sizeof(1)) in body
    assert "int" in body
    assert a.address in body
    assert "memory" in body

@gen_cluster(
    client=True,
    scheduler_kwargs={"dashboard": True},
    config={
        "distributed.scheduler.dashboard.bokeh-application.allow_websocket_origin": [
            "good.invalid"
        ]
    },
)
async def test_allow_websocket_origin(c, s, a, b):
    from tornado.httpclient import HTTPRequest
    from tornado.websocket import websocket_connect

    url = (
        "ws://localhost:%d/status/ws?bokeh-protocol-version=1.0&bokeh-session-id=1"
        % s.http_server.port
    )
    with pytest.raises(HTTPClientError) as err:
        await websocket_connect(
            HTTPRequest(url, headers={"Origin": "http://evil.invalid"})
        )
    assert err.value.code == 403

@gen_cluster(client=True)
async def test_eventstream(c, s, a, b):
    from tornado.websocket import websocket_connect

    ws_client = await websocket_connect(
        "ws://localhost:%d/%s" % (s.http_server.port, "eventstream")
    )
    assert "websocket" in str(s.plugins).lower()
    ws_client.close()
<commit_msg>Remove `pytest.mark.repeat` from `test_prometheus_collect_task_states` (#5376)<commit_after>import asyncio
import json
import re

import pytest

pytest.importorskip("bokeh")

from tornado.escape import url_escape
from tornado.httpclient import AsyncHTTPClient, HTTPClientError

from dask.sizeof import sizeof

from distributed.utils import is_valid_xml
from distributed.utils_test import gen_cluster, inc, slowinc

@gen_cluster(client=True)
async def test_connect(c, s, a, b):
    future = c.submit(lambda x: x + 1, 1)
    x = c.submit(slowinc, 1, delay=1, retries=5)
    await future
    http_client = AsyncHTTPClient()
    for suffix in [
        "info/main/workers.html",
        "info/worker/" + url_escape(a.address) + ".html",
        "info/task/" + url_escape(future.key) + ".html",
        "info/main/logs.html",
        "info/logs/" + url_escape(a.address) + ".html",
        "info/call-stack/" + url_escape(x.key) + ".html",
        "info/call-stacks/" + url_escape(a.address) + ".html",
        "json/counts.json",
        "json/identity.json",
        "json/index.html",
        "individual-plots.json",
        "sitemap.json",
    ]:
        response = await http_client.fetch(
            "http://localhost:%d/%s" % (s.http_server.port, suffix)
        )
        assert response.code == 200
        body = response.body.decode()
        if suffix.endswith(".json"):
            json.loads(body)
        else:
            assert is_valid_xml(body)
            assert not re.search("href=./", body)  # no absolute links

@gen_cluster(client=True, nthreads=[])
async def test_worker_404(c, s):
    http_client = AsyncHTTPClient()
    with pytest.raises(HTTPClientError) as err:
        await http_client.fetch(
            "http://localhost:%d/info/worker/unknown" % s.http_server.port
        )
    assert err.value.code == 404
    with pytest.raises(HTTPClientError) as err:
        await http_client.fetch(
            "http://localhost:%d/info/task/unknown" % s.http_server.port
        )
    assert err.value.code == 404

@gen_cluster(client=True, scheduler_kwargs={"http_prefix": "/foo", "dashboard": True})
async def test_prefix(c, s, a, b):
    http_client = AsyncHTTPClient()
    for suffix in ["foo/info/main/workers.html", "foo/json/index.html", "foo/system"]:
        response = await http_client.fetch(
            "http://localhost:%d/%s" % (s.http_server.port, suffix)
        )
        assert response.code == 200
        body = response.body.decode()
        if suffix.endswith(".json"):
            json.loads(body)
        else:
            assert is_valid_xml(body)

@gen_cluster(client=True, clean_kwargs={"threads": False})
async def test_prometheus(c, s, a, b):
    pytest.importorskip("prometheus_client")
    from prometheus_client.parser import text_string_to_metric_families

    http_client = AsyncHTTPClient()

    # request data twice since there once was a case where metrics got registered multiple times resulting in
    # prometheus_client errors
    for _ in range(2):
        response = await http_client.fetch(
            "http://localhost:%d/metrics" % s.http_server.port
        )
        assert response.code == 200
        assert response.headers["Content-Type"] == "text/plain; version=0.0.4"

        txt = response.body.decode("utf8")
        families = {
            family.name: family for family in text_string_to_metric_families(txt)
        }
        assert "dask_scheduler_workers" in families

        client = families["dask_scheduler_clients"]
        assert client.samples[0].value == 1.0

@gen_cluster(client=True, clean_kwargs={"threads": False})
async def test_prometheus_collect_task_states(c, s, a, b):
    pytest.importorskip("prometheus_client")
    from prometheus_client.parser import text_string_to_metric_families

    http_client = AsyncHTTPClient()

    async def fetch_metrics():
        port = s.http_server.port
        response = await http_client.fetch(f"http://localhost:{port}/metrics")
        txt = response.body.decode("utf8")
        families = {
            family.name: family for family in text_string_to_metric_families(txt)
        }

        active_metrics = {
            sample.labels["state"]: sample.value
            for sample in families["dask_scheduler_tasks"].samples
        }
        forgotten_tasks = [
            sample.value
            for sample in families["dask_scheduler_tasks_forgotten"].samples
        ]
        return active_metrics, forgotten_tasks

    expected = {"memory", "released", "processing", "waiting", "no-worker", "erred"}

    # Ensure that we get full zero metrics for all states even though the
    # scheduler did nothing, yet
    assert not s.tasks
    active_metrics, forgotten_tasks = await fetch_metrics()
    assert active_metrics.keys() == expected
    assert sum(active_metrics.values()) == 0.0
    assert sum(forgotten_tasks) == 0.0

    # submit a task which should show up in the prometheus scraping
    future = c.submit(slowinc, 1, delay=0.5)
    while not any(future.key in w.tasks for w in [a, b]):
        await asyncio.sleep(0.001)

    active_metrics, forgotten_tasks = await fetch_metrics()
    assert active_metrics.keys() == expected
    assert sum(active_metrics.values()) == 1.0
    assert sum(forgotten_tasks) == 0.0

    res = await c.gather(future)
    assert res == 2

    future.release()

    while any(future.key in w.tasks for w in [a, b]):
        await asyncio.sleep(0.001)

    active_metrics, forgotten_tasks = await fetch_metrics()
    assert active_metrics.keys() == expected
    assert sum(active_metrics.values()) == 0.0
    assert sum(forgotten_tasks) == 0.0

@gen_cluster(client=True, clean_kwargs={"threads": False})
async def test_health(c, s, a, b):
    http_client = AsyncHTTPClient()

    response = await http_client.fetch(
        "http://localhost:%d/health" % s.http_server.port
    )
    assert response.code == 200
    assert response.headers["Content-Type"] == "text/plain"

    txt = response.body.decode("utf8")
    assert txt == "ok"

@gen_cluster()
async def test_sitemap(s, a, b):
    http_client = AsyncHTTPClient()

    response = await http_client.fetch(
        "http://localhost:%d/sitemap.json" % s.http_server.port
    )
    out = json.loads(response.body.decode())
    assert "paths" in out
    assert "/sitemap.json" in out["paths"]
    assert "/health" in out["paths"]
    assert "/statics/css/base.css" in out["paths"]

@gen_cluster(client=True)
async def test_task_page(c, s, a, b):
    future = c.submit(lambda x: x + 1, 1, workers=a.address)
    x = c.submit(inc, 1)
    await future
    http_client = AsyncHTTPClient()

    "info/task/" + url_escape(future.key) + ".html",
    response = await http_client.fetch(
        "http://localhost:%d/info/task/" % s.http_server.port
        + url_escape(future.key)
        + ".html"
    )
    assert response.code == 200
    body = response.body.decode()

    assert str(sizeof(1)) in body
    assert "int" in body
    assert a.address in body
    assert "memory" in body

@gen_cluster(
    client=True,
    scheduler_kwargs={"dashboard": True},
    config={
        "distributed.scheduler.dashboard.bokeh-application.allow_websocket_origin": [
            "good.invalid"
        ]
    },
)
async def test_allow_websocket_origin(c, s, a, b):
    from tornado.httpclient import HTTPRequest
    from tornado.websocket import websocket_connect

    url = (
        "ws://localhost:%d/status/ws?bokeh-protocol-version=1.0&bokeh-session-id=1"
        % s.http_server.port
    )
    with pytest.raises(HTTPClientError) as err:
        await websocket_connect(
            HTTPRequest(url, headers={"Origin": "http://evil.invalid"})
        )
    assert err.value.code == 403

@gen_cluster(client=True)
async def test_eventstream(c, s, a, b):
    from tornado.websocket import websocket_connect

    ws_client = await websocket_connect(
        "ws://localhost:%d/%s" % (s.http_server.port, "eventstream")
    )
    assert "websocket" in str(s.plugins).lower()
    ws_client.close()
<|endoftext|>