<commit_before>from os.path import expanduser
import random
import socket
from contextlib import closing
import os
import threading
import common
import inspect

COORDINATOR_NAME = "coordinator"
WORKER1 = "worker1"
WORKER2 = "worker2"

REGULAR_USER_NAME = "regularuser"
SUPER_USER_NAME = "postgres"

ARBITRARY_SCHEDULE_NAMES = [
    "create_schedule",
    "sql_schedule",
    "sql_base_schedule",
    "postgres_schedule",
]

BEFORE_PG_UPGRADE_SCHEDULE = "./before_pg_upgrade_schedule"
AFTER_PG_UPGRADE_SCHEDULE = "./after_pg_upgrade_schedule"

CREATE_SCHEDULE = "./create_schedule"
POSTGRES_SCHEDULE = "./postgres_schedule"
SQL_SCHEDULE = "./sql_schedule"
SQL_BASE_SCHEDULE = "./sql_base_schedule"

AFTER_CITUS_UPGRADE_COORD_SCHEDULE = "./after_citus_upgrade_coord_schedule"
BEFORE_CITUS_UPGRADE_COORD_SCHEDULE = "./before_citus_upgrade_coord_schedule"
MIXED_BEFORE_CITUS_UPGRADE_SCHEDULE = "./mixed_before_citus_upgrade_schedule"
MIXED_AFTER_CITUS_UPGRADE_SCHEDULE = "./mixed_after_citus_upgrade_schedule"

CITUS_ARBITRARY_TEST_DIR = "./tmp_citus_test"

MASTER = "master"
# This should be updated when citus version changes
MASTER_VERSION = "11.0"

HOME = expanduser("~")

CITUS_VERSION_SQL = "SELECT extversion FROM pg_extension WHERE extname = 'citus';"

# this is out of ephemeral port range for many systems hence
# it is a lower change that it will conflict with "in-use" ports
next_port = 10200

# ephemeral port start on many linux systems
PORT_UPPER = 32768

port_lock = threading.Lock()

def should_include_config(class_name):

    if inspect.isclass(class_name) and (
        issubclass(class_name, CitusMXBaseClusterConfig)
        or issubclass(class_name, CitusDefaultClusterConfig)
    ):
        return True
    return False

def find_free_port():
    global next_port
    with port_lock:
        while next_port < PORT_UPPER:
            with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
                try:
                    s.bind(("localhost", next_port))
                    port = next_port
                    next_port += 1
                    return port
                except:
                    next_port += 1
    # we couldn't find a port
    raise Exception("Couldn't find a port to use")

class NewInitCaller(type):
    def __call__(cls, *args, **kwargs):
        obj = type.__call__(cls, *args, **kwargs)
        obj.post_init()
        return obj

class CitusBaseClusterConfig(object, metaclass=NewInitCaller):
    def __init__(self, arguments):
        if "--bindir" in arguments:
            self.bindir = arguments["--bindir"]
        self.pg_srcdir = arguments["--pgxsdir"]
        self.temp_dir = CITUS_ARBITRARY_TEST_DIR
        self.worker_amount = 2
        self.user = REGULAR_USER_NAME
        self.is_mx = False
        self.is_citus = True
        self.name = type(self).__name__
        self.settings = {
            "shared_preload_libraries": "citus",
            "log_error_verbosity": "terse",
            "fsync": False,
            "citus.node_conninfo": "sslmode=prefer",
            "citus.enable_repartition_joins": True,
            "citus.repartition_join_bucket_count_per_node": 2,
            "max_connections": 600,
        }
        self.new_settings = {}
        self.add_coordinator_to_metadata = False
        self.env_variables = {}

    def post_init(self):
        self._init_node_name_ports()

        self.datadir = self.temp_dir + "/data"
        self.datadir += self.name
        self.input_dir = self.datadir
        self.output_dir = self.datadir
        self.output_file = os.path.join(self.datadir, "run.out")
        if self.worker_amount > 0:
            self.chosen_random_worker_port = self.random_worker_port()
        self.settings.update(self.new_settings)

    def coordinator_port(self):
        return self.node_name_to_ports[COORDINATOR_NAME]

    def setup_steps(self):
        pass

    def random_worker_port(self):
        return random.choice(self.worker_ports)

    def random_port(self):
        return random.choice(list(self.node_name_to_ports.values()))

    def _init_node_name_ports(self):
        self.node_name_to_ports = {}
        self.worker_ports = []
        cur_port = self._get_and_update_next_port()
        self.node_name_to_ports[COORDINATOR_NAME] = cur_port
        for i in range(self.worker_amount):
            cur_port = self._get_and_update_next_port()
            cur_worker_name = "worker{}".format(i)
            self.node_name_to_ports[cur_worker_name] = cur_port
            self.worker_ports.append(cur_port)

    def _get_and_update_next_port(self):
        if hasattr(self, "fixed_port"):
            next_port = self.fixed_port
            self.fixed_port += 1
            return next_port
        return find_free_port()

class CitusDefaultClusterConfig(CitusBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        new_settings = {
            "client_min_messages": "WARNING",
            "citus.sort_returning": True,
            "citus.use_citus_managed_tables": True,
        }
        self.settings.update(new_settings)
        self.add_coordinator_to_metadata = True

class CitusMXBaseClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.is_mx = True

class CitusUpgradeConfig(CitusBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.pre_tar_path = arguments["--citus-pre-tar"]
        self.post_tar_path = arguments["--citus-post-tar"]
        self.temp_dir = "./tmp_citus_upgrade"
        self.new_settings = {"citus.enable_version_checks": "false"}
        self.user = SUPER_USER_NAME
        self.mixed_mode = arguments["--mixed"]
        self.fixed_port = 57635

class PostgresConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.worker_amount = 0
        self.is_citus = False
        self.new_settings = {
            "citus.use_citus_managed_tables": False,
        }

class CitusSingleNodeClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.worker_amount = 0

    def setup_steps(self):
        common.coordinator_should_haveshards(self.bindir, self.coordinator_port())

class CitusSingleWorkerClusterConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.worker_amount = 1

class CitusSuperUserDefaultClusterConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.user = SUPER_USER_NAME

class CitusThreeWorkersManyShardsClusterConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {"citus.shard_count": 191}
        self.worker_amount = 3

    def setup_steps(self):
        common.coordinator_should_haveshards(self.bindir, self.coordinator_port())

class CitusSmallSharedPoolSizeConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
             "citus.local_shared_pool_size": 5,
            "citus.max_shared_pool_size": 5,
        }

class CitusSmallExecutorPoolSizeConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.max_adaptive_executor_pool_size": 2,
        }

class CitusSequentialExecutionConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.multi_shard_modify_mode": "sequential",
        }

class CitusCacheManyConnectionsConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.max_cached_conns_per_worker": 4,
        }

class CitusUnusualExecutorConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.max_adaptive_executor_pool_size": 7,
            "citus.executor_slow_start_interval": 1,
            "citus.prevent_incomplete_connection_establishment": False,
            "citus.enable_cost_based_connection_establishment": False,
            "citus.max_cached_connection_lifetime": "10ms",
            # https://github.com/citusdata/citus/issues/5345
            # "citus.force_max_query_parallelization": "on",
            "citus.binary_worker_copy_format": False,
            "citus.enable_binary_protocol": False,
            "citus.local_table_join_policy": "prefer-distributed",
        }

        # this setting does not necessarily need to be here
        # could go any other test
        self.env_variables = {'PGAPPNAME' : 'test_app'}

class CitusSmallCopyBuffersConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.copy_switchover_threshold": "1B",
            "citus.local_copy_flush_threshold": "1B",
            "citus.remote_copy_flush_threshold": "1B",
        }

class CitusUnusualQuerySettingsConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.use_citus_managed_tables": False,
            "citus.task_assignment_policy": "first-replica",
            "citus.enable_fast_path_router_planner": False,
            "citus.enable_local_execution": False,
            "citus.enable_single_hash_repartition_joins": True,
            "citus.recover_2pc_interval": "1s",
            "citus.remote_task_check_interval": "1ms",
            "citus.values_materialization_threshold": "0",
        }

class CitusSingleNodeSingleShardClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.worker_amount = 0
        self.new_settings = {"citus.shard_count": 1}

    def setup_steps(self):
        common.coordinator_should_haveshards(self.bindir, self.coordinator_port())

class CitusShardReplicationFactorClusterConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {"citus.shard_replication_factor": 2}

class CitusSingleShardClusterConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {"citus.shard_count": 1}

class CitusNonMxClusterConfig(CitusMXBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.is_mx = False

class PGUpgradeConfig(CitusBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.old_bindir = arguments["--old-bindir"]
        self.new_bindir = arguments["--new-bindir"]
        self.temp_dir = "./tmp_upgrade"
        self.old_datadir = self.temp_dir + "/oldData"
        self.new_datadir = self.temp_dir + "/newData"
        self.user = SUPER_USER_NAME
        self.add_coordinator_to_metadata = True
<commit_msg>Turn metadata sync on in arbitrary config tests<commit_after>from os.path import expanduser
import random
import socket
from contextlib import closing
import os
import threading
import common
import inspect

COORDINATOR_NAME = "coordinator"
WORKER1 = "worker1"
WORKER2 = "worker2"

REGULAR_USER_NAME = "regularuser"
SUPER_USER_NAME = "postgres"

ARBITRARY_SCHEDULE_NAMES = [
    "create_schedule",
    "sql_schedule",
    "sql_base_schedule",
    "postgres_schedule",
]

BEFORE_PG_UPGRADE_SCHEDULE = "./before_pg_upgrade_schedule"
AFTER_PG_UPGRADE_SCHEDULE = "./after_pg_upgrade_schedule"

CREATE_SCHEDULE = "./create_schedule"
POSTGRES_SCHEDULE = "./postgres_schedule"
SQL_SCHEDULE = "./sql_schedule"
SQL_BASE_SCHEDULE = "./sql_base_schedule"

AFTER_CITUS_UPGRADE_COORD_SCHEDULE = "./after_citus_upgrade_coord_schedule"
BEFORE_CITUS_UPGRADE_COORD_SCHEDULE = "./before_citus_upgrade_coord_schedule"
MIXED_BEFORE_CITUS_UPGRADE_SCHEDULE = "./mixed_before_citus_upgrade_schedule"
MIXED_AFTER_CITUS_UPGRADE_SCHEDULE = "./mixed_after_citus_upgrade_schedule"

CITUS_ARBITRARY_TEST_DIR = "./tmp_citus_test"

MASTER = "master"
# This should be updated when citus version changes
MASTER_VERSION = "11.0"

HOME = expanduser("~")

CITUS_VERSION_SQL = "SELECT extversion FROM pg_extension WHERE extname = 'citus';"

# this is out of ephemeral port range for many systems hence
# it is a lower change that it will conflict with "in-use" ports
next_port = 10200

# ephemeral port start on many linux systems
PORT_UPPER = 32768

port_lock = threading.Lock()

def should_include_config(class_name):

    if inspect.isclass(class_name) and issubclass(class_name, CitusDefaultClusterConfig):
        return True
    return False

def find_free_port():
    global next_port
    with port_lock:
        while next_port < PORT_UPPER:
            with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
                try:
                    s.bind(("localhost", next_port))
                    port = next_port
                    next_port += 1
                    return port
                except:
                    next_port += 1
    # we couldn't find a port
    raise Exception("Couldn't find a port to use")

class NewInitCaller(type):
    def __call__(cls, *args, **kwargs):
        obj = type.__call__(cls, *args, **kwargs)
        obj.post_init()
        return obj

class CitusBaseClusterConfig(object, metaclass=NewInitCaller):
    def __init__(self, arguments):
        if "--bindir" in arguments:
            self.bindir = arguments["--bindir"]
        self.pg_srcdir = arguments["--pgxsdir"]
        self.temp_dir = CITUS_ARBITRARY_TEST_DIR
        self.worker_amount = 2
        self.user = REGULAR_USER_NAME
        self.is_mx = True
        self.is_citus = True
        self.name = type(self).__name__
        self.settings = {
            "shared_preload_libraries": "citus",
            "log_error_verbosity": "terse",
            "fsync": False,
            "citus.node_conninfo": "sslmode=prefer",
            "citus.enable_repartition_joins": True,
            "citus.repartition_join_bucket_count_per_node": 2,
            "max_connections": 600,
        }
        self.new_settings = {}
        self.add_coordinator_to_metadata = False
        self.env_variables = {}

    def post_init(self):
        self._init_node_name_ports()

        self.datadir = self.temp_dir + "/data"
        self.datadir += self.name
        self.input_dir = self.datadir
        self.output_dir = self.datadir
        self.output_file = os.path.join(self.datadir, "run.out")
        if self.worker_amount > 0:
            self.chosen_random_worker_port = self.random_worker_port()
        self.settings.update(self.new_settings)

    def coordinator_port(self):
        return self.node_name_to_ports[COORDINATOR_NAME]

    def setup_steps(self):
        pass

    def random_worker_port(self):
        return random.choice(self.worker_ports)

    def random_port(self):
        return random.choice(list(self.node_name_to_ports.values()))

    def _init_node_name_ports(self):
        self.node_name_to_ports = {}
        self.worker_ports = []
        cur_port = self._get_and_update_next_port()
        self.node_name_to_ports[COORDINATOR_NAME] = cur_port
        for i in range(self.worker_amount):
            cur_port = self._get_and_update_next_port()
            cur_worker_name = "worker{}".format(i)
            self.node_name_to_ports[cur_worker_name] = cur_port
            self.worker_ports.append(cur_port)

    def _get_and_update_next_port(self):
        if hasattr(self, "fixed_port"):
            next_port = self.fixed_port
            self.fixed_port += 1
            return next_port
        return find_free_port()

class CitusDefaultClusterConfig(CitusBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        new_settings = {
            "client_min_messages": "WARNING",
            "citus.sort_returning": True,
            "citus.use_citus_managed_tables": True,
        }
        self.settings.update(new_settings)
        self.add_coordinator_to_metadata = True

class CitusUpgradeConfig(CitusBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.pre_tar_path = arguments["--citus-pre-tar"]
        self.post_tar_path = arguments["--citus-post-tar"]
        self.temp_dir = "./tmp_citus_upgrade"
        self.new_settings = {"citus.enable_version_checks": "false"}
        self.user = SUPER_USER_NAME
        self.mixed_mode = arguments["--mixed"]
        self.fixed_port = 57635
        self.is_mx = False

class PostgresConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.worker_amount = 0
        self.is_citus = False
        self.new_settings = {
            "citus.use_citus_managed_tables": False,
        }

class CitusSingleNodeClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.worker_amount = 0

    def setup_steps(self):
        common.coordinator_should_haveshards(self.bindir, self.coordinator_port())

class CitusSingleWorkerClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.worker_amount = 1

class CitusSuperUserDefaultClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.user = SUPER_USER_NAME

class CitusThreeWorkersManyShardsClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {"citus.shard_count": 191}
        self.worker_amount = 3

    def setup_steps(self):
        common.coordinator_should_haveshards(self.bindir, self.coordinator_port())

class CitusSmallSharedPoolSizeConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
             "citus.local_shared_pool_size": 5,
            "citus.max_shared_pool_size": 5,
        }

class CitusSmallExecutorPoolSizeConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.max_adaptive_executor_pool_size": 2,
        }

class CitusSequentialExecutionConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.multi_shard_modify_mode": "sequential",
        }

class CitusCacheManyConnectionsConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.max_cached_conns_per_worker": 4,
        }

class CitusUnusualExecutorConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.max_adaptive_executor_pool_size": 7,
            "citus.executor_slow_start_interval": 1,
            "citus.prevent_incomplete_connection_establishment": False,
            "citus.enable_cost_based_connection_establishment": False,
            "citus.max_cached_connection_lifetime": "10ms",
            # https://github.com/citusdata/citus/issues/5345
            # "citus.force_max_query_parallelization": "on",
            "citus.binary_worker_copy_format": False,
            "citus.enable_binary_protocol": False,
            "citus.local_table_join_policy": "prefer-distributed",
        }

        # this setting does not necessarily need to be here
        # could go any other test
        self.env_variables = {'PGAPPNAME' : 'test_app'}

class CitusSmallCopyBuffersConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.copy_switchover_threshold": "1B",
            "citus.local_copy_flush_threshold": "1B",
            "citus.remote_copy_flush_threshold": "1B",
        }

class CitusUnusualQuerySettingsConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {
            "citus.use_citus_managed_tables": False,
            "citus.task_assignment_policy": "first-replica",
            "citus.enable_fast_path_router_planner": False,
            "citus.enable_local_execution": False,
            "citus.enable_single_hash_repartition_joins": True,
            "citus.recover_2pc_interval": "1s",
            "citus.remote_task_check_interval": "1ms",
            "citus.values_materialization_threshold": "0",
        }

class CitusSingleNodeSingleShardClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.worker_amount = 0
        self.new_settings = {"citus.shard_count": 1}

    def setup_steps(self):
        common.coordinator_should_haveshards(self.bindir, self.coordinator_port())

class CitusShardReplicationFactorClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {"citus.shard_replication_factor": 2}

class CitusSingleShardClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.new_settings = {"citus.shard_count": 1}

class CitusNonMxClusterConfig(CitusDefaultClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.is_mx = False

class PGUpgradeConfig(CitusBaseClusterConfig):
    def __init__(self, arguments):
        super().__init__(arguments)
        self.old_bindir = arguments["--old-bindir"]
        self.new_bindir = arguments["--new-bindir"]
        self.temp_dir = "./tmp_upgrade"
        self.old_datadir = self.temp_dir + "/oldData"
        self.new_datadir = self.temp_dir + "/newData"
        self.user = SUPER_USER_NAME
        self.add_coordinator_to_metadata = True
<|endoftext|>


<commit_before>from __future__ import print_function, absolute_import, division

from unittest2 import TestCase
from mock import patch, call
import os
import signal
import subprocess
import tempfile

from succubus import Daemon

class TestDaemonize(TestCase):
    @patch('succubus.daemonize.sys')
    def test_must_pop_sys_argv_before_loading_config(self, mock_sys):
        """The sys.argv.pop() must happen before load_configuration()

        This way, load_configuration() has a chance to parse the command
        line arguments, which may contain something like a --config=xyz
        parameter that affects config loading.
        """
        class MyDaemon(Daemon):
            def load_configuration(self):
                if self.param1 != 'start':
                    raise Exception("param1 not yet set")

        mock_sys.argv = ['foo', 'start', '--config=xyz']

        MyDaemon(pid_file='foo.pid')

    @patch("succubus.daemonize.os.setgid")
    def test_set_gid_translates_group_name(self, mock_setgid):
        daemon = Daemon(pid_file="foo")
        daemon.group = "root"

        daemon.set_gid()

        mock_setgid.assert_called_with(0)

    @patch("succubus.daemonize.os.setuid")
    def test_set_uid_translates_user_name(self, mock_setuid):
        daemon = Daemon(pid_file="foo")
        daemon.user = "root"

        daemon.set_uid()

        mock_setuid.assert_called_with(0)

    @patch("succubus.daemonize.sys")
    def test_daemon_insists_on_pidfile(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']
        self.assertRaisesRegexp(Exception, "You did not provide a pid file",
                                Daemon)

    @patch("succubus.daemonize.sys")
    def test_daemon_remembers_abspath_of_pidfile(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']
        daemon = Daemon(pid_file="foo")
        self.assertTrue(os.path.isabs(daemon.pid_file))

    @patch("succubus.daemonize.sys")
    def test_delpid_really_deletes_pidfile(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']

        fake_pidfile = tempfile.NamedTemporaryFile()
        try:
            daemon = Daemon(pid_file=fake_pidfile.name)
            daemon.delpid()
            self.assertFalse(os.path.exists(fake_pidfile.name))
        finally:
            if os.path.exists(fake_pidfile.name):
                os.unlink(fake_pidfile.name)

    @patch("succubus.daemonize.sys")
    def test_already_running_detects_running_process(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']

        fake_pidfile = tempfile.NamedTemporaryFile()
        try:
            fake_pidfile.write("{pid}\n".format(pid=os.getpid()))
            fake_pidfile.flush()

            daemon = Daemon(pid_file=fake_pidfile.name)
            self.assertEqual(daemon._already_running(), True)
            # _already_running must set the self.pid.
            self.assertEqual(daemon.pid, os.getpid())
        finally:
            os.unlink(fake_pidfile.name)

    @patch("succubus.daemonize.sys")
    def test_already_running_handles_missing_pidfile(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']

        daemon = Daemon(pid_file="/a_nonexisting_pidfile")
        self.assertEqual(daemon._already_running(), False)

    @patch("succubus.daemonize.sys")
    def test_already_running_handles_missing_process(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']

        # The test assumes that the OS will not reuse test_pid before
        # this function is finished.
        test_process = subprocess.Popen(['true'])
        test_pid = test_process.pid
        test_process.wait()

        fake_pidfile = tempfile.NamedTemporaryFile()
        try:
            fake_pidfile.write("{pid}\n".format(pid=test_pid))
            fake_pidfile.flush()

            daemon = Daemon(pid_file=fake_pidfile.name)
            self.assertEqual(daemon._already_running(), False)
        finally:
            os.unlink(fake_pidfile.name)

    @patch("succubus.daemonize.time.sleep")
    @patch("succubus.daemonize.os.kill")
    @patch("succubus.daemonize.sys")
    def test_reliable_kill_actually_kills(self, mock_sys, mock_kill, mock_sleep):
        mock_sys.argv = ['foo', 'bar']
        daemon = Daemon(pid_file="foo")
        daemon.pid = 1234

        # Pretend the second os.kill() fails because the process terminated.
        mock_kill.side_effect = [None, OSError("No such process")]

        retval = daemon.reliable_kill()

        self.assertEqual(retval, 0)
        self.assertEqual(mock_kill.call_count, 2)
        self.assertEqual(mock_sleep.call_count, 1)
        expected_call = call(1234, signal.SIGTERM)
        mock_kill.assert_has_calls([expected_call, expected_call])

    @patch("succubus.daemonize.time.sleep")
    @patch("succubus.daemonize.os.kill")
    @patch("succubus.daemonize.sys")
    def test_reliable_kill_kills_stuck_processes(self, mock_sys, mock_kill, mock_sleep):
        """When a process ignores SIGTERM, it must be killed with SIGKILL"""
        mock_sys.argv = ['foo', 'bar']
        daemon = Daemon(pid_file="foo")
        daemon.pid = 1234

        mock_kill.side_effect = None

        retval = daemon.reliable_kill()

        self.assertEqual(retval, 0)
        # Must have tried 100 times with SIGTERM and 1 time with SIGKILL.
        self.assertEqual(mock_kill.call_count, 101)
        mock_kill.assert_any_call(1234, signal.SIGKILL)
<commit_msg>Fix test problems in Python 3<commit_after>from __future__ import print_function, absolute_import, division

from unittest2 import TestCase
from mock import patch, call
import os
import signal
import six
import subprocess
import tempfile

from succubus import Daemon

class TestDaemonize(TestCase):
    @patch('succubus.daemonize.sys')
    def test_must_pop_sys_argv_before_loading_config(self, mock_sys):
        """The sys.argv.pop() must happen before load_configuration()

        This way, load_configuration() has a chance to parse the command
        line arguments, which may contain something like a --config=xyz
        parameter that affects config loading.
        """
        class MyDaemon(Daemon):
            def load_configuration(self):
                if self.param1 != 'start':
                    raise Exception("param1 not yet set")

        mock_sys.argv = ['foo', 'start', '--config=xyz']

        MyDaemon(pid_file='foo.pid')

    @patch("succubus.daemonize.os.setgid")
    def test_set_gid_translates_group_name(self, mock_setgid):
        daemon = Daemon(pid_file="foo")
        daemon.group = "root"

        daemon.set_gid()

        mock_setgid.assert_called_with(0)

    @patch("succubus.daemonize.os.setuid")
    def test_set_uid_translates_user_name(self, mock_setuid):
        daemon = Daemon(pid_file="foo")
        daemon.user = "root"

        daemon.set_uid()

        mock_setuid.assert_called_with(0)

    @patch("succubus.daemonize.sys")
    def test_daemon_insists_on_pidfile(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']
        self.assertRaisesRegexp(Exception, "You did not provide a pid file",
                                Daemon)

    @patch("succubus.daemonize.sys")
    def test_daemon_remembers_abspath_of_pidfile(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']
        daemon = Daemon(pid_file="foo")
        self.assertTrue(os.path.isabs(daemon.pid_file))

    @patch("succubus.daemonize.sys")
    def test_delpid_really_deletes_pidfile(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']

        fake_pidfile = tempfile.NamedTemporaryFile()
        try:
            daemon = Daemon(pid_file=fake_pidfile.name)
            daemon.delpid()
            self.assertFalse(os.path.exists(fake_pidfile.name))
        finally:
            if os.path.exists(fake_pidfile.name):
                os.unlink(fake_pidfile.name)

    @patch("succubus.daemonize.sys")
    def test_already_running_detects_running_process(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']

        fake_pidfile = tempfile.NamedTemporaryFile()
        try:
            content = "{pid}\n".format(pid=os.getpid())
            fake_pidfile.write(six.b(content))
            fake_pidfile.flush()

            daemon = Daemon(pid_file=fake_pidfile.name)
            self.assertEqual(daemon._already_running(), True)
            # _already_running must set the self.pid.
            self.assertEqual(daemon.pid, os.getpid())
        finally:
            os.unlink(fake_pidfile.name)

    @patch("succubus.daemonize.sys")
    def test_already_running_handles_missing_pidfile(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']

        daemon = Daemon(pid_file="/a_nonexisting_pidfile")
        self.assertEqual(daemon._already_running(), False)

    @patch("succubus.daemonize.sys")
    def test_already_running_handles_missing_process(self, mock_sys):
        mock_sys.argv = ['foo', 'bar']

        # The test assumes that the OS will not reuse test_pid before
        # this function is finished.
        test_process = subprocess.Popen(['true'])
        test_pid = test_process.pid
        test_process.wait()

        fake_pidfile = tempfile.NamedTemporaryFile()
        try:
            content = "{pid}\n".format(pid=test_pid)
            fake_pidfile.write(six.b(content))
            fake_pidfile.flush()

            daemon = Daemon(pid_file=fake_pidfile.name)
            self.assertEqual(daemon._already_running(), False)
        finally:
            os.unlink(fake_pidfile.name)

    @patch("succubus.daemonize.time.sleep")
    @patch("succubus.daemonize.os.kill")
    @patch("succubus.daemonize.sys")
    def test_reliable_kill_actually_kills(self, mock_sys, mock_kill, mock_sleep):
        mock_sys.argv = ['foo', 'bar']
        daemon = Daemon(pid_file="foo")
        daemon.pid = 1234

        # Pretend the second os.kill() fails because the process terminated.
        mock_kill.side_effect = [None, OSError("No such process")]

        retval = daemon.reliable_kill()

        self.assertEqual(retval, 0)
        self.assertEqual(mock_kill.call_count, 2)
        self.assertEqual(mock_sleep.call_count, 1)
        expected_call = call(1234, signal.SIGTERM)
        mock_kill.assert_has_calls([expected_call, expected_call])

    @patch("succubus.daemonize.time.sleep")
    @patch("succubus.daemonize.os.kill")
    @patch("succubus.daemonize.sys")
    def test_reliable_kill_kills_stuck_processes(self, mock_sys, mock_kill, mock_sleep):
        """When a process ignores SIGTERM, it must be killed with SIGKILL"""
        mock_sys.argv = ['foo', 'bar']
        daemon = Daemon(pid_file="foo")
        daemon.pid = 1234

        mock_kill.side_effect = None

        retval = daemon.reliable_kill()

        self.assertEqual(retval, 0)
        # Must have tried 100 times with SIGTERM and 1 time with SIGKILL.
        self.assertEqual(mock_kill.call_count, 101)
        mock_kill.assert_any_call(1234, signal.SIGKILL)
<|endoftext|>


<commit_before>from __future__ import with_statement
import re
import os
import sys
from urllib import quote, unquote
from time import gmtime, strftime, mktime, strptime, time

MANIFEST_DIR = "manifests"
_re_client = re.compile("client-(?P<lang>[^.]+)\.xml")
_re_resource = re.compile("(?:<script[^>]*?src=\"([^\"]*))|(?:<link[^>]*?href=\"([^\"]*))")

def get_timestamp(path = None):
    return strftime("%a, %d %b %Y %H:%M:%S GMT", gmtime(path and stat(path).st_mtime or None))

def get_resources(os_path, web_path, file_name):
    resources = []
    with open(os.path.join(os_path, file_name), 'r') as f:
        content = f.read()
        for match in _re_resource.finditer(content):
            resource = filter(bool, match.groups())[0]
            if resource.startswith("/"):
                pass
            elif resource.startswith("./"):
                resource = web_path + resource[1:]
            elif resource.startswidth("../"):
                while resource.startswith("../"):
                    pos = resource.find("/", 3)
                    if pos > -1:
                        resource = resource[pos:]
                    else:
                        resource = ""
            else:
                resourc = web_path + "/" + resource
            resources.append(resource)
    return resources

def write_maifest(os_path, file_name, resources, tag=""):
    with open(os.path.join(os_path, MANIFEST_DIR, file_name), "wb") as f:
        content = ["CACHE MANIFEST", "# created %s %s" % (get_timestamp(), tag)]
        content.extend(resources)
        content.append("")
        f.write("\n".join(content))

def add_manifest(os_path, web_path, client_file, manifest):
    content = None
    manifest_path = "/".join([web_path, MANIFEST_DIR, manifest])
    with open(os.path.join(os_path, client_file), "rb") as f:
        content = f.read()
    if content and not ' manifest="' in content:
        with open(os.path.join(os_path, client_file), "wb") as f:
            f.write(content.replace(">", ' manifest="%s">' % manifest_path, 1))

def main(argv=sys.argv):
    """

    """
    import optparse
    usage = """%prog [options] [path]

    Create a manifest for all files in path with the pattern client-<lang code>.xml.
    """
    parser = optparse.OptionParser(usage)
    parser.add_option("-d", "--domain", type="string", dest="domain",
                      default=None, help="The domain token of the.")
    parser.add_option("-t", "--tag", dest="tag", default="",
                      help="A tag to be used in the manifest comment.")
    options, args = parser.parse_args()
    os_path = os.path.abspath(len(args) > 0 and args[0] or ".")
    web_path = os_path
    if options.domain:
        pos = web_path.find(options.domain)
        if pos > -1:
            web_path = web_path[pos + len(options.domain):]
    web_path = "/".join(map(lambda p: quote(p), web_path.split(os.path.sep)))
    if not os.path.exists(os.path.join(os_path, MANIFEST_DIR)):
        os.makedirs(os.path.join(os_path, MANIFEST_DIR))
    for name in os.listdir(os_path):
         match = _re_client.search(name)
         if match:
            resources = [web_path + "/", web_path + "/" + name]
            manifest = "dragonfly-%s.manifest" % match.group("lang")
            if web_path.startswith("/app"):
                if "cutting-edge" in web_path:
                    resources.append("/app/cutting-edge/")
                else:
                    resources.append("/app/")
            resources.extend(get_resources(os_path, web_path, name))
            write_maifest(os_path, manifest, resources, options.tag)
            add_manifest(os_path, web_path, name, manifest)

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Reverted reg exp. Simplified (and fixed) resource paths.<commit_after>from __future__ import with_statement
import re
import os
import sys
from urllib import quote, unquote
from time import gmtime, strftime, mktime, strptime, time

MANIFEST_DIR = "manifests"
_re_client = re.compile("client-(?P<lang>[^.]+)\.xml")
_re_resource = re.compile("(?:<script[^>]*src=\"([^\"]*))|(?:<link[^>]*href=\"([^\"]*))")

def get_timestamp(path = None):
    return strftime("%a, %d %b %Y %H:%M:%S GMT", gmtime(path and stat(path).st_mtime or None))

def get_resources(os_path, web_path, file_name):
    resources = []
    with open(os.path.join(os_path, file_name), 'r') as f:
        content = f.read()
        for match in _re_resource.finditer(content):
            resource = filter(bool, match.groups())[0]
            if resource.startswith("/"):
                pass
            elif resource.startswith("./"):
                resource = web_path + resource[1:]
            # everything else resolved by the browser
            else:
                resourc = web_path + "/" + resource
            resources.append(resource)
    return resources

def write_maifest(os_path, file_name, resources, tag=""):
    with open(os.path.join(os_path, MANIFEST_DIR, file_name), "wb") as f:
        content = ["CACHE MANIFEST", "# created %s %s" % (get_timestamp(), tag)]
        content.extend(resources)
        content.append("")
        f.write("\n".join(content))

def add_manifest(os_path, web_path, client_file, manifest):
    content = None
    manifest_path = "/".join([web_path, MANIFEST_DIR, manifest])
    with open(os.path.join(os_path, client_file), "rb") as f:
        content = f.read()
    if content and not ' manifest="' in content:
        with open(os.path.join(os_path, client_file), "wb") as f:
            f.write(content.replace(">", ' manifest="%s">' % manifest_path, 1))

def main(argv=sys.argv):
    """

    """
    import optparse
    usage = """%prog [options] [path]

    Create a manifest for all files in path with the pattern client-<lang code>.xml.
    """
    parser = optparse.OptionParser(usage)
    parser.add_option("-d", "--domain", type="string", dest="domain",
                      default=None, help="The domain token of the.")
    parser.add_option("-t", "--tag", dest="tag", default="",
                      help="A tag to be used in the manifest comment.")
    options, args = parser.parse_args()
    os_path = os.path.abspath(len(args) > 0 and args[0] or ".")
    web_path = os_path
    if options.domain:
        pos = web_path.find(options.domain)
        if pos > -1:
            web_path = web_path[pos + len(options.domain):]
    web_path = "/".join(map(lambda p: quote(p), web_path.split(os.path.sep)))
    if not os.path.exists(os.path.join(os_path, MANIFEST_DIR)):
        os.makedirs(os.path.join(os_path, MANIFEST_DIR))
    for name in os.listdir(os_path):
         match = _re_client.search(name)
         if match:
            resources = [web_path + "/", web_path + "/" + name]
            manifest = "dragonfly-%s.manifest" % match.group("lang")
            if web_path.startswith("/app"):
                if "cutting-edge" in web_path:
                    resources.append("/app/cutting-edge/")
                else:
                    resources.append("/app/")
            resources.extend(get_resources(os_path, web_path, name))
            write_maifest(os_path, manifest, resources, options.tag)
            add_manifest(os_path, web_path, name, manifest)

if __name__ == "__main__":
    sys.exit(main())
<|endoftext|>


<commit_before># Creates a Zulip remote development environment droplet or
# a production droplet in DigitalOcean.
#
# Particularly useful for sprints/hackathons, interns, and other
# situation where one wants to quickly onboard new contributors.
#
# Requires python-digitalocean library:
# https://github.com/koalalorenzo/python-digitalocean
#
# Also requires DigitalOcean team membership for Zulip and API token:
# https://cloud.digitalocean.com/settings/api/tokens
#
# Copy conf.ini-template to conf.ini and populate with your API token.
import argparse
import configparser
import json
import os
import sys
import time
import urllib.error
import urllib.request
from typing import Any, Dict, List

import digitalocean
import requests

parser = argparse.ArgumentParser(description="Create a Zulip development VM DigitalOcean droplet.")
parser.add_argument(
    "username", help="GitHub username for whom you want to create a Zulip dev droplet"
)
parser.add_argument("--tags", nargs="+", default=[])
parser.add_argument("-f", "--recreate", action="store_true")
parser.add_argument("-p", "--production", action="store_true")

def get_config() -> configparser.ConfigParser:
    config = configparser.ConfigParser()
    config.read(os.path.join(os.path.dirname(os.path.abspath(__file__)), "conf.ini"))
    return config

def assert_github_user_exists(github_username: str) -> bool:
    print(f"Checking to see if GitHub user {github_username} exists...")
    user_api_url = f"https://api.github.com/users/{github_username}"
    try:
        response = urllib.request.urlopen(user_api_url)
        json.load(response)
        print("...user exists!")
        return True
    except urllib.error.HTTPError as err:
        print(err)
        print(f"Does the GitHub user {github_username} exist?")
        sys.exit(1)

def get_ssh_public_keys_from_github(github_username: str) -> List[Dict[str, Any]]:
    print("Checking to see that GitHub user has available public keys...")
    apiurl_keys = f"https://api.github.com/users/{github_username}/keys"
    try:
        response = urllib.request.urlopen(apiurl_keys)
        userkeys = json.load(response)
        if not userkeys:
            print(
                f"No keys found. Has user {github_username} added SSH keys to their GitHub account?"
            )
            sys.exit(1)
        print("...public keys found!")
        return userkeys
    except urllib.error.HTTPError as err:
        print(err)
        print(f"Has user {github_username} added SSH keys to their GitHub account?")
        sys.exit(1)

def assert_user_forked_zulip_server_repo(username: str) -> bool:
    print("Checking to see GitHub user has forked zulip/zulip...")
    apiurl_fork = f"https://api.github.com/repos/{username}/zulip"
    try:
        response = urllib.request.urlopen(apiurl_fork)
        json.load(response)
        print("...fork found!")
        return True
    except urllib.error.HTTPError as err:
        print(err)
        print(f"Has user {username} forked zulip/zulip?")
        sys.exit(1)

def assert_droplet_does_not_exist(my_token: str, droplet_name: str, recreate: bool) -> None:
    print(f"Checking to see if droplet {droplet_name} already exists...")
    manager = digitalocean.Manager(token=my_token)
    my_droplets = manager.get_all_droplets()
    for droplet in my_droplets:
        if droplet.name.lower() == droplet_name:
            if not recreate:
                print(
                    "Droplet {} already exists. Pass --recreate if you "
                    "need to recreate the droplet.".format(droplet_name)
                )
                sys.exit(1)
            else:
                print(f"Deleting existing droplet {droplet_name}.")
                droplet.destroy()
                return
    print("...No droplet found...proceeding.")

def get_ssh_keys_string_from_github_ssh_key_dicts(userkey_dicts: List[Dict[str, Any]]) -> str:
    return "\n".join([userkey_dict["key"] for userkey_dict in userkey_dicts])

def generate_dev_droplet_user_data(username: str, userkey_dicts: List[Dict[str, Any]]) -> str:
    ssh_keys_string = get_ssh_keys_string_from_github_ssh_key_dicts(userkey_dicts)
    setup_root_ssh_keys = f"printf '{ssh_keys_string}' > /root/.ssh/authorized_keys"
    setup_zulipdev_ssh_keys = f"printf '{ssh_keys_string}' > /home/zulipdev/.ssh/authorized_keys"

    # We pass the hostname as username.zulipdev.org to the DigitalOcean API.
    # But some droplets (eg on 18.04) are created with with hostname set to just username.
    # So we fix the hostname using cloud-init.
    hostname_setup = f"hostnamectl set-hostname {username}.zulipdev.org"

    setup_repo = (
        "cd /home/zulipdev/{1} && "
        "git remote add origin https://github.com/{0}/{1}.git && "
        "git fetch origin && "
        "git clean -f"
    )

    server_repo_setup = setup_repo.format(username, "zulip")
    python_api_repo_setup = setup_repo.format(username, "python-zulip-api")

    cloudconf = f"""\
#!/bin/bash

{setup_zulipdev_ssh_keys}
{setup_root_ssh_keys}
sed -i "s/PasswordAuthentication yes/PasswordAuthentication no/g" /etc/ssh/sshd_config
service ssh restart
{hostname_setup}
su -c '{server_repo_setup}' zulipdev
su -c '{python_api_repo_setup}' zulipdev
su -c 'git config --global core.editor nano' zulipdev
su -c 'git config --global pull.rebase true' zulipdev
"""
    print("...returning cloud-config data.")
    return cloudconf

def generate_prod_droplet_user_data(username: str, userkey_dicts: List[Dict[str, Any]]) -> str:
    ssh_keys_string = get_ssh_keys_string_from_github_ssh_key_dicts(userkey_dicts)
    setup_root_ssh_keys = f"printf '{ssh_keys_string}' > /root/.ssh/authorized_keys"

    cloudconf = f"""\
#!/bin/bash

{setup_root_ssh_keys}
passwd -d root
sed -i "s/PasswordAuthentication yes/PasswordAuthentication no/g" /etc/ssh/sshd_config
service ssh restart
"""
    print("...returning cloud-config data.")
    return cloudconf

def create_droplet(
    my_token: str, template_id: str, name: str, tags: List[str], user_data: str
) -> str:
    droplet = digitalocean.Droplet(
        token=my_token,
        name=name,
        region="nyc3",
        image=template_id,
        size_slug="s-1vcpu-2gb",
        user_data=user_data,
        tags=tags,
        backups=False,
    )

    print("Initiating droplet creation...")
    droplet.create()

    incomplete = True
    while incomplete:
        actions = droplet.get_actions()
        for action in actions:
            action.load()
            print(f"...[{action.type}]: {action.status}")
            if action.type == "create" and action.status == "completed":
                incomplete = False
                break
        if incomplete:
            time.sleep(15)
    print("...droplet created!")
    droplet.load()
    print(f"...ip address for new droplet is: {droplet.ip_address}.")
    return droplet.ip_address

def delete_existing_records(records: List[digitalocean.Record], record_name: str) -> None:
    count = 0
    for record in records:
        if record.name == record_name and record.domain == "zulipdev.org" and record.type == "A":
            record.destroy()
            count = count + 1
    if count:
        print(f"Deleted {count} existing A records for {record_name}.zulipdev.org.")

def create_dns_record(my_token: str, record_name: str, ip_address: str) -> None:
    domain = digitalocean.Domain(token=my_token, name="zulipdev.org")
    domain.load()
    records = domain.get_records()

    delete_existing_records(records, record_name)
    wildcard_name = "*." + record_name
    delete_existing_records(records, wildcard_name)

    print(f"Creating new A record for {record_name}.zulipdev.org that points to {ip_address}.")
    domain.create_new_domain_record(type="A", name=record_name, data=ip_address)
    print(f"Creating new A record for *.{record_name}.zulipdev.org that points to {ip_address}.")
    domain.create_new_domain_record(type="A", name=wildcard_name, data=ip_address)

def print_dev_droplet_instructions(droplet_domain_name: str) -> None:
    print(
        """
COMPLETE! Droplet for GitHub user {0} is available at {0}.zulipdev.org.

Instructions for use are below. (copy and paste to the user)

------
Your remote Zulip dev server has been created!

- Connect to your server by running
  `ssh zulipdev@{0}` on the command line
  (Terminal for macOS and Linux, Bash for Git on Windows).
- There is no password; your account is configured to use your SSH keys.
- Once you log in, you should see `(zulip-py3-venv) ~$`.
- To start the dev server, `cd zulip` and then run `./tools/run-dev.py`.
- While the dev server is running, you can see the Zulip server in your browser at
  http://{0}:9991.
""".format(
            droplet_domain_name
        )
    )

    print(
        "See [Developing remotely](https://zulip.readthedocs.io/en/latest/development/remote.html) "
        "for tips on using the remote dev instance and "
        "[Git & GitHub guide](https://zulip.readthedocs.io/en/latest/git/index.html) "
        "to learn how to use Git with Zulip.\n"
    )
    print(
        "Note that this droplet will automatically be deleted after a month of inactivity. "
        "If you are leaving Zulip for more than a few weeks, we recommend pushing all of your "
        "active branches to GitHub."
    )
    print("------")

def print_production_droplet_instructions(droplet_domain_name: str) -> None:
    print(
        """
-----

Production droplet created successfully!

Connect to the server by running

ssh root@{}

-----
""".format(
            droplet_domain_name
        )
    )

def get_zulip_oneclick_app_slug(api_token: str) -> str:
    response = requests.get(
        "https://api.digitalocean.com/v2/1-clicks", headers={"Authorization": f"Bearer {api_token}"}
    ).json()
    one_clicks = response["1_clicks"]

    for one_click in one_clicks:
        if one_click["slug"].startswith("kandralabs"):
            return one_click["slug"]
    raise Exception("Unable to find Zulip One-click app slug")

if __name__ == "__main__":
    args = parser.parse_args()
    username = args.username.lower()

    if args.production:
        print(f"Creating production droplet for GitHub user {username}...")
    else:
        print(f"Creating Zulip developer environment for GitHub user {username}...")

    config = get_config()
    api_token = config["digitalocean"]["api_token"]

    assert_github_user_exists(github_username=username)

    public_keys = get_ssh_public_keys_from_github(github_username=username)

    if args.production:
        subdomain = f"{username}-prod"
        droplet_domain_name = f"{subdomain}.zulipdev.org"
        template_id = get_zulip_oneclick_app_slug(api_token)
        user_data = generate_prod_droplet_user_data(username=username, userkey_dicts=public_keys)

    else:
        assert_user_forked_zulip_server_repo(username=username)

        subdomain = username
        droplet_domain_name = f"{subdomain}.zulipdev.org"
        user_data = generate_dev_droplet_user_data(username=username, userkey_dicts=public_keys)

        # define id of image to create new droplets from
        # You can get this with something like the following. You may need to try other pages.
        # Broken in two to satisfy linter (line too long)
        # curl -X GET -H "Content-Type: application/json" -u <API_KEY>: "https://api.digitaloc
        # ean.com/v2/images?page=5" | grep --color=always base.zulipdev.org
        template_id = "63219191"

    assert_droplet_does_not_exist(
        my_token=api_token, droplet_name=droplet_domain_name, recreate=args.recreate
    )

    ip_address = create_droplet(
        my_token=api_token,
        template_id=template_id,
        name=droplet_domain_name,
        tags=args.tags,
        user_data=user_data,
    )

    create_dns_record(my_token=api_token, record_name=subdomain, ip_address=ip_address)

    if args.production:
        print_production_droplet_instructions(droplet_domain_name=droplet_domain_name)
    else:
        print_dev_droplet_instructions(droplet_domain_name=droplet_domain_name)

    sys.exit(1)
<commit_msg>droplet: Allow overriding the subdomain.<commit_after># Creates a Zulip remote development environment droplet or
# a production droplet in DigitalOcean.
#
# Particularly useful for sprints/hackathons, interns, and other
# situation where one wants to quickly onboard new contributors.
#
# Requires python-digitalocean library:
# https://github.com/koalalorenzo/python-digitalocean
#
# Also requires DigitalOcean team membership for Zulip and API token:
# https://cloud.digitalocean.com/settings/api/tokens
#
# Copy conf.ini-template to conf.ini and populate with your API token.
import argparse
import configparser
import json
import os
import sys
import time
import urllib.error
import urllib.request
from typing import Any, Dict, List

import digitalocean
import requests

parser = argparse.ArgumentParser(description="Create a Zulip development VM DigitalOcean droplet.")
parser.add_argument(
    "username", help="GitHub username for whom you want to create a Zulip dev droplet"
)
parser.add_argument("--tags", nargs="+", default=[])
parser.add_argument("-f", "--recreate", action="store_true")
parser.add_argument("-s", "--subdomain")
parser.add_argument("-p", "--production", action="store_true")

def get_config() -> configparser.ConfigParser:
    config = configparser.ConfigParser()
    config.read(os.path.join(os.path.dirname(os.path.abspath(__file__)), "conf.ini"))
    return config

def assert_github_user_exists(github_username: str) -> bool:
    print(f"Checking to see if GitHub user {github_username} exists...")
    user_api_url = f"https://api.github.com/users/{github_username}"
    try:
        response = urllib.request.urlopen(user_api_url)
        json.load(response)
        print("...user exists!")
        return True
    except urllib.error.HTTPError as err:
        print(err)
        print(f"Does the GitHub user {github_username} exist?")
        sys.exit(1)

def get_ssh_public_keys_from_github(github_username: str) -> List[Dict[str, Any]]:
    print("Checking to see that GitHub user has available public keys...")
    apiurl_keys = f"https://api.github.com/users/{github_username}/keys"
    try:
        response = urllib.request.urlopen(apiurl_keys)
        userkeys = json.load(response)
        if not userkeys:
            print(
                f"No keys found. Has user {github_username} added SSH keys to their GitHub account?"
            )
            sys.exit(1)
        print("...public keys found!")
        return userkeys
    except urllib.error.HTTPError as err:
        print(err)
        print(f"Has user {github_username} added SSH keys to their GitHub account?")
        sys.exit(1)

def assert_user_forked_zulip_server_repo(username: str) -> bool:
    print("Checking to see GitHub user has forked zulip/zulip...")
    apiurl_fork = f"https://api.github.com/repos/{username}/zulip"
    try:
        response = urllib.request.urlopen(apiurl_fork)
        json.load(response)
        print("...fork found!")
        return True
    except urllib.error.HTTPError as err:
        print(err)
        print(f"Has user {username} forked zulip/zulip?")
        sys.exit(1)

def assert_droplet_does_not_exist(my_token: str, droplet_name: str, recreate: bool) -> None:
    print(f"Checking to see if droplet {droplet_name} already exists...")
    manager = digitalocean.Manager(token=my_token)
    my_droplets = manager.get_all_droplets()
    for droplet in my_droplets:
        if droplet.name.lower() == droplet_name:
            if not recreate:
                print(
                    "Droplet {} already exists. Pass --recreate if you "
                    "need to recreate the droplet.".format(droplet_name)
                )
                sys.exit(1)
            else:
                print(f"Deleting existing droplet {droplet_name}.")
                droplet.destroy()
                return
    print("...No droplet found...proceeding.")

def get_ssh_keys_string_from_github_ssh_key_dicts(userkey_dicts: List[Dict[str, Any]]) -> str:
    return "\n".join([userkey_dict["key"] for userkey_dict in userkey_dicts])

def generate_dev_droplet_user_data(
    username: str, subdomain: str, userkey_dicts: List[Dict[str, Any]]
) -> str:
    ssh_keys_string = get_ssh_keys_string_from_github_ssh_key_dicts(userkey_dicts)
    setup_root_ssh_keys = f"printf '{ssh_keys_string}' > /root/.ssh/authorized_keys"
    setup_zulipdev_ssh_keys = f"printf '{ssh_keys_string}' > /home/zulipdev/.ssh/authorized_keys"

    # We pass the hostname as username.zulipdev.org to the DigitalOcean API.
    # But some droplets (eg on 18.04) are created with with hostname set to just username.
    # So we fix the hostname using cloud-init.
    hostname_setup = f"hostnamectl set-hostname {subdomain}.zulipdev.org"

    setup_repo = (
        "cd /home/zulipdev/{1} && "
        "git remote add origin https://github.com/{0}/{1}.git && "
        "git fetch origin && "
        "git clean -f"
    )

    server_repo_setup = setup_repo.format(username, "zulip")
    python_api_repo_setup = setup_repo.format(username, "python-zulip-api")

    cloudconf = f"""\
#!/bin/bash

{setup_zulipdev_ssh_keys}
{setup_root_ssh_keys}
sed -i "s/PasswordAuthentication yes/PasswordAuthentication no/g" /etc/ssh/sshd_config
service ssh restart
{hostname_setup}
su -c '{server_repo_setup}' zulipdev
su -c '{python_api_repo_setup}' zulipdev
su -c 'git config --global core.editor nano' zulipdev
su -c 'git config --global pull.rebase true' zulipdev
"""
    print("...returning cloud-config data.")
    return cloudconf

def generate_prod_droplet_user_data(username: str, userkey_dicts: List[Dict[str, Any]]) -> str:
    ssh_keys_string = get_ssh_keys_string_from_github_ssh_key_dicts(userkey_dicts)
    setup_root_ssh_keys = f"printf '{ssh_keys_string}' > /root/.ssh/authorized_keys"

    cloudconf = f"""\
#!/bin/bash

{setup_root_ssh_keys}
passwd -d root
sed -i "s/PasswordAuthentication yes/PasswordAuthentication no/g" /etc/ssh/sshd_config
service ssh restart
"""
    print("...returning cloud-config data.")
    return cloudconf

def create_droplet(
    my_token: str, template_id: str, name: str, tags: List[str], user_data: str
) -> str:
    droplet = digitalocean.Droplet(
        token=my_token,
        name=name,
        region="nyc3",
        image=template_id,
        size_slug="s-1vcpu-2gb",
        user_data=user_data,
        tags=tags,
        backups=False,
    )

    print("Initiating droplet creation...")
    droplet.create()

    incomplete = True
    while incomplete:
        actions = droplet.get_actions()
        for action in actions:
            action.load()
            print(f"...[{action.type}]: {action.status}")
            if action.type == "create" and action.status == "completed":
                incomplete = False
                break
        if incomplete:
            time.sleep(15)
    print("...droplet created!")
    droplet.load()
    print(f"...ip address for new droplet is: {droplet.ip_address}.")
    return droplet.ip_address

def delete_existing_records(records: List[digitalocean.Record], record_name: str) -> None:
    count = 0
    for record in records:
        if record.name == record_name and record.domain == "zulipdev.org" and record.type == "A":
            record.destroy()
            count = count + 1
    if count:
        print(f"Deleted {count} existing A records for {record_name}.zulipdev.org.")

def create_dns_record(my_token: str, record_name: str, ip_address: str) -> None:
    domain = digitalocean.Domain(token=my_token, name="zulipdev.org")
    domain.load()
    records = domain.get_records()

    delete_existing_records(records, record_name)
    wildcard_name = "*." + record_name
    delete_existing_records(records, wildcard_name)

    print(f"Creating new A record for {record_name}.zulipdev.org that points to {ip_address}.")
    domain.create_new_domain_record(type="A", name=record_name, data=ip_address)
    print(f"Creating new A record for *.{record_name}.zulipdev.org that points to {ip_address}.")
    domain.create_new_domain_record(type="A", name=wildcard_name, data=ip_address)

def print_dev_droplet_instructions(droplet_domain_name: str) -> None:
    print(
        """
COMPLETE! Droplet for GitHub user {0} is available at {0}.zulipdev.org.

Instructions for use are below. (copy and paste to the user)

------
Your remote Zulip dev server has been created!

- Connect to your server by running
  `ssh zulipdev@{0}` on the command line
  (Terminal for macOS and Linux, Bash for Git on Windows).
- There is no password; your account is configured to use your SSH keys.
- Once you log in, you should see `(zulip-py3-venv) ~$`.
- To start the dev server, `cd zulip` and then run `./tools/run-dev.py`.
- While the dev server is running, you can see the Zulip server in your browser at
  http://{0}:9991.
""".format(
            droplet_domain_name
        )
    )

    print(
        "See [Developing remotely](https://zulip.readthedocs.io/en/latest/development/remote.html) "
        "for tips on using the remote dev instance and "
        "[Git & GitHub guide](https://zulip.readthedocs.io/en/latest/git/index.html) "
        "to learn how to use Git with Zulip.\n"
    )
    print(
        "Note that this droplet will automatically be deleted after a month of inactivity. "
        "If you are leaving Zulip for more than a few weeks, we recommend pushing all of your "
        "active branches to GitHub."
    )
    print("------")

def print_production_droplet_instructions(droplet_domain_name: str) -> None:
    print(
        """
-----

Production droplet created successfully!

Connect to the server by running

ssh root@{}

-----
""".format(
            droplet_domain_name
        )
    )

def get_zulip_oneclick_app_slug(api_token: str) -> str:
    response = requests.get(
        "https://api.digitalocean.com/v2/1-clicks", headers={"Authorization": f"Bearer {api_token}"}
    ).json()
    one_clicks = response["1_clicks"]

    for one_click in one_clicks:
        if one_click["slug"].startswith("kandralabs"):
            return one_click["slug"]
    raise Exception("Unable to find Zulip One-click app slug")

if __name__ == "__main__":
    args = parser.parse_args()
    username = args.username.lower()
    if args.subdomain:
        subdomain = args.subdomain.lower()
    elif args.production:
        subdomain = "{username}-prod"
    else:
        subdomain = username

    if args.production:
        print(f"Creating production droplet for GitHub user {username}...")
    else:
        print(f"Creating Zulip developer environment for GitHub user {username}...")

    config = get_config()
    api_token = config["digitalocean"]["api_token"]

    assert_github_user_exists(github_username=username)

    public_keys = get_ssh_public_keys_from_github(github_username=username)

    if args.production:
        droplet_domain_name = f"{subdomain}.zulipdev.org"
        template_id = get_zulip_oneclick_app_slug(api_token)
        user_data = generate_prod_droplet_user_data(username=username, userkey_dicts=public_keys)

    else:
        assert_user_forked_zulip_server_repo(username=username)

        droplet_domain_name = f"{subdomain}.zulipdev.org"
        user_data = generate_dev_droplet_user_data(
            username=username, subdomain=subdomain, userkey_dicts=public_keys
        )

        # define id of image to create new droplets from
        # You can get this with something like the following. You may need to try other pages.
        # Broken in two to satisfy linter (line too long)
        # curl -X GET -H "Content-Type: application/json" -u <API_KEY>: "https://api.digitaloc
        # ean.com/v2/images?page=5" | grep --color=always base.zulipdev.org
        template_id = "63219191"

    assert_droplet_does_not_exist(
        my_token=api_token, droplet_name=droplet_domain_name, recreate=args.recreate
    )

    ip_address = create_droplet(
        my_token=api_token,
        template_id=template_id,
        name=droplet_domain_name,
        tags=args.tags,
        user_data=user_data,
    )

    create_dns_record(my_token=api_token, record_name=subdomain, ip_address=ip_address)

    if args.production:
        print_production_droplet_instructions(droplet_domain_name=droplet_domain_name)
    else:
        print_dev_droplet_instructions(droplet_domain_name=droplet_domain_name)

    sys.exit(1)
<|endoftext|>


<commit_before>#!/usr/bin/env python

#  Copyright 2008-2009 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""Robot Framework Test Data Documentation Tool

Usage:  testdoc.py [options] data_sources

This tool generates a high level test documentation from a given test data.
Generated documentation includes the names, documentations and other metadata
of each test suite and test case, as well as the top-level keywords and their
arguments. Most of the options accepted by this tool have exactly same
semantics as same options have when executing test cases.

Options:
  -o --output path       Where to write the generated documentation. If the
                         path is a directory, the documentation is
                         generated there using name '<suitename>-doc.html'.
  -T --title title       Set the title of the generated documentation.
                         Underscores in the title are converted to spaces.
  -N --name name         Set the name of the top level test suite.
  -D --doc document      Set the document of the top level test suite.
  -M --metadata name:value *  Set metadata of the top level test suite.
  -G --settag tag *      Set given tag(s) to all test cases.
  -t --test name *       Include test cases by name.
  -s --suite name *      Include test suites by name.
  -i --include tag *     Include test cases by tags.
  -e --exclude tag *     Exclude test cases by tags.
  -h --help              Print this help.

Examples:
  $ testdoc.py mytestcases.html
  $ testdoc.py --name smoke_test_plan --include smoke path/to/my_tests/
"""

import sys
import os
import time

from robot import utils
from robot.common import BaseKeyword, BaseTestSuite
from robot.running import TestSuite, Keyword
from robot.conf import RobotSettings
from robot.serializing.logserializers import LogSerializer
from robot.serializing import templates
from robot.serializing.templating import Namespace, Template
from robot.errors import DataError, Information
from robot.parsing import rawdata
from robot.variables import Variables

rawdata.PROCESS_CURDIR = False
Variables.set_from_variable_table = lambda self, varz: None

def generate_test_doc(args):
    opts, datasources = process_arguments(args)
    suite = TestSuite(datasources, RobotSettings(opts))
    outpath = get_outpath(opts['output'], suite.name)
    serialize_test_doc(suite, outpath, opts['title'])
    exit(msg=outpath)

def serialize_test_doc(suite, outpath, title):
    outfile = open(outpath, 'w')
    serializer = TestdocSerializer(outfile, suite)
    ttuple = time.localtime()
    str_time = utils.format_time(ttuple, daytimesep=' ', gmtsep=' ')
    int_time = long(time.mktime(ttuple))
    if title:
        title = title.replace('_', ' ')
    else:
        title = 'Documentation for %s' % suite.name
    namespace = Namespace(gentime_str=str_time, gentime_int=int_time,
                          version=utils.get_full_version('testdoc.py'),
                          suite=suite, title=title)
    Template(template=templates.LOG).generate(namespace, outfile)
    suite.serialize(serializer)
    outfile.write('\n\n')
    outfile.close()

def process_arguments(args_list):
    argparser = utils.ArgumentParser(__doc__)
    try:
        opts, args = argparser.parse_args(args_list, help='help', check_args=True)
    except Information, msg:
        exit(msg=str(msg))
    except DataError, err:
        exit(error=str(err))
    return opts, args

def exit(msg=None, error=None):
    if msg:
        sys.stdout.write(msg + '\n')
        sys.exit(0)
    sys.stderr.write(error + '\n\nTry --help for usage information.\n')
    sys.exit(1)

def get_outpath(path, suite_name):
    if not path:
        path = '.'
    if os.path.isdir(path):
        path = os.path.join(path, '%s-doc.html' % suite_name.replace(' ', '_'))
    return os.path.abspath(path)

class TestdocSerializer(LogSerializer):

    def __init__(self, output, suite):
        self._writer = utils.HtmlWriter(output)
        self._idgen = utils.IdGenerator()
        self._split_level = -1
        self._suite_level = 0

    def start_suite(self, suite):
        suite._set_variable_dependent_metadata(NonResolvingContext())
        LogSerializer.start_suite(self, suite)

    def start_test(self, test):
        test._init_test(NonResolvingContext())
        LogSerializer.start_test(self, test)

    def start_keyword(self, kw):
        if isinstance(kw, Keyword):  # Doesn't match For or Parallel
            kw.name = kw._get_name(kw.name, NonResolvingContext())
        LogSerializer.start_keyword(self, kw)

    def _is_element_open(self, item):
        return isinstance(item, BaseTestSuite)

    def _write_times(self, item):
        pass

    def _write_suite_metadata(self, suite):
        self._start_suite_or_test_metadata(suite)
        for name, value in suite.get_metadata(html=True):
            self._write_metadata_row(name, value, escape=False, write_empty=True)
        self._write_source(suite.source)
        self._write_metadata_row('Number of Tests', suite.get_test_count())
        self._writer.end('table')

    def _write_test_metadata(self, test):
        self._start_suite_or_test_metadata(test)
        if test.timeout.secs < 0:
            tout = ''
        else:
            tout = utils.secs_to_timestr(test.timeout.secs)
            if test.timeout.message:
                tout += ' | ' + test.timeout.message
        self._write_metadata_row('Timeout', tout)
        self._write_metadata_row('Tags', ', '.join(test.tags))
        self._writer.end('table')

    def _write_folding_button(self, item):
        if not isinstance(item, BaseKeyword):
            LogSerializer._write_folding_button(self, item)

    def _write_expand_all(self, item):
        if isinstance(item, BaseTestSuite):
            LogSerializer._write_expand_all(self, item)

class NonResolvingContext:

    def replace_vars_from_setting(self, name, item, errors):
        return item

if __name__ == '__main__':
    generate_test_doc(sys.argv[1:])
<commit_msg>Fixed after API changes<commit_after>#!/usr/bin/env python

#  Copyright 2008-2009 Nokia Siemens Networks Oyj
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""Robot Framework Test Data Documentation Tool

Usage:  testdoc.py [options] data_sources

This tool generates a high level test documentation from a given test data.
Generated documentation includes the names, documentations and other metadata
of each test suite and test case, as well as the top-level keywords and their
arguments. Most of the options accepted by this tool have exactly same
semantics as same options have when executing test cases.

Options:
  -o --output path       Where to write the generated documentation. If the
                         path is a directory, the documentation is
                         generated there using name '<suitename>-doc.html'.
  -T --title title       Set the title of the generated documentation.
                         Underscores in the title are converted to spaces.
  -N --name name         Set the name of the top level test suite.
  -D --doc document      Set the document of the top level test suite.
  -M --metadata name:value *  Set metadata of the top level test suite.
  -G --settag tag *      Set given tag(s) to all test cases.
  -t --test name *       Include test cases by name.
  -s --suite name *      Include test suites by name.
  -i --include tag *     Include test cases by tags.
  -e --exclude tag *     Exclude test cases by tags.
  -h --help              Print this help.

Examples:
  $ testdoc.py mytestcases.html
  $ testdoc.py --name smoke_test_plan --include smoke path/to/my_tests/
"""

import sys
import os
import time

from robot import utils
from robot.common import BaseKeyword, BaseTestSuite
from robot.running import TestSuite, Keyword
from robot.conf import RobotSettings
from robot.serializing.logserializers import LogSerializer
from robot.serializing import templates
from robot.serializing.templating import Namespace, Template
from robot.errors import DataError, Information
from robot.parsing import datareader
from robot.variables import Variables

datareader.PROCESS_CURDIR = False
Variables.set_from_variable_table = lambda self, varz: None

def generate_test_doc(args):
    opts, datasources = process_arguments(args)
    suite = TestSuite(datasources, RobotSettings(opts))
    outpath = get_outpath(opts['output'], suite.name)
    serialize_test_doc(suite, outpath, opts['title'])
    exit(msg=outpath)

def serialize_test_doc(suite, outpath, title):
    outfile = open(outpath, 'w')
    serializer = TestdocSerializer(outfile, suite)
    ttuple = time.localtime()
    str_time = utils.format_time(ttuple, daytimesep=' ', gmtsep=' ')
    int_time = long(time.mktime(ttuple))
    if title:
        title = title.replace('_', ' ')
    else:
        title = 'Documentation for %s' % suite.name
    namespace = Namespace(gentime_str=str_time, gentime_int=int_time,
                          version=utils.get_full_version('testdoc.py'),
                          suite=suite, title=title)
    Template(template=templates.LOG).generate(namespace, outfile)
    suite.serialize(serializer)
    outfile.write('\n\n')
    outfile.close()

def process_arguments(args_list):
    argparser = utils.ArgumentParser(__doc__)
    try:
        opts, args = argparser.parse_args(args_list, help='help', check_args=True)
    except Information, msg:
        exit(msg=str(msg))
    except DataError, err:
        exit(error=str(err))
    return opts, args

def exit(msg=None, error=None):
    if msg:
        sys.stdout.write(msg + '\n')
        sys.exit(0)
    sys.stderr.write(error + '\n\nTry --help for usage information.\n')
    sys.exit(1)

def get_outpath(path, suite_name):
    if not path:
        path = '.'
    if os.path.isdir(path):
        path = os.path.join(path, '%s-doc.html' % suite_name.replace(' ', '_'))
    return os.path.abspath(path)

class TestdocSerializer(LogSerializer):

    def __init__(self, output, suite):
        self._writer = utils.HtmlWriter(output)
        self._idgen = utils.IdGenerator()
        self._split_level = -1
        self._suite_level = 0

    def start_suite(self, suite):
        suite._set_variable_dependent_metadata(NonResolvingContext())
        LogSerializer.start_suite(self, suite)

    def start_test(self, test):
        test._init_test(NonResolvingContext())
        LogSerializer.start_test(self, test)

    def start_keyword(self, kw):
        if isinstance(kw, Keyword):  # Doesn't match For
            kw.name = kw._get_name(kw.name)
        LogSerializer.start_keyword(self, kw)

    def _is_element_open(self, item):
        return isinstance(item, BaseTestSuite)

    def _write_times(self, item):
        pass

    def _write_suite_metadata(self, suite):
        self._start_suite_or_test_metadata(suite)
        for name, value in suite.get_metadata(html=True):
            self._write_metadata_row(name, value, escape=False, write_empty=True)
        self._write_source(suite.source)
        self._write_metadata_row('Number of Tests', suite.get_test_count())
        self._writer.end('table')

    def _write_test_metadata(self, test):
        self._start_suite_or_test_metadata(test)
        if test.timeout.secs < 0:
            tout = ''
        else:
            tout = utils.secs_to_timestr(test.timeout.secs)
            if test.timeout.message:
                tout += ' | ' + test.timeout.message
        self._write_metadata_row('Timeout', tout)
        self._write_metadata_row('Tags', ', '.join(test.tags))
        self._writer.end('table')

    def _write_folding_button(self, item):
        if not isinstance(item, BaseKeyword):
            LogSerializer._write_folding_button(self, item)

    def _write_expand_all(self, item):
        if isinstance(item, BaseTestSuite):
            LogSerializer._write_expand_all(self, item)

class NonResolvingContext:

    def replace_vars_from_setting(self, name, item, errors):
        return item

    def replace_string(self, item):
        return item

    def get_current_vars(self):
        return NonResolvingContext()

if __name__ == '__main__':
    generate_test_doc(sys.argv[1:])
<|endoftext|>


<commit_before># Copyright 2014-2015 Isotoma Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import time
try:
    from contextlib import ExitStack
except ImportError:
    from contextlib2 import ExitStack

from touchdown.core.action import Action
from touchdown.core.resource import Resource
from touchdown.core.plan import Plan
from touchdown.core import argument, errors

from ..account import BaseAccount
from ..common import SimpleDescribe, SimpleApply, SimpleDestroy

class Image(Resource):

    resource_name = "image"

    name = argument.String(field="Name")
    description = argument.String(field="Description")

    ami = argument.String()

    resources = argument.List()

    #architecture = argument.String(field="Architecture", default="x86_64", choices=["x86_64", "i386"])
    #kernel = argument.String(field="KernelId")
    #ramdisk = argument.String(field="RamdiskId")
    #root_device_name = argument.String(field="RootDeviceName")
    #virtualization_type = argument.String(choices=["paravirtual", "hvm"], field="VirtualizationType")
    #sriov_net_support = argument.String(choices=["simple"], field="SriovNetSupport")
    #location = argument.String()
    #snapshot_id = argument.String()

    account = argument.Resource(BaseAccount)

class BuildInstance(Action):

    @property
    def description(self):
        yield "Build new AMI '{}' from '{}'".format(self.resource.name, self.resource.ami)

    def create_security_group(self):
        print("Creating temporary security group")
        security_group = self.plan.client.create_security_group(
            GroupName="temporary-security-group",
            Description="Temporary security group",
        )
        self.stack.callback(self.destroy_security_group, security_group)
        return security_group

    def destroy_security_group(self, security_group):
        print("Deleting temproary security group")
        self.plan.client.delete_security_group(
            GroupId=security_group["GroupId"],
        )

    def create_keypair(self):
        print("Creating temporary keypair")
        keypair = self.plan.client.create_key_pair(
            KeyName="temporary-key-pair",
        )
        self.stack.callback(self.destroy_keypair, keypair)
        return keypair

    def destroy_keypair(self, keypair):
        print("Deleting temproary keypair")
        self.plan.client.delete_key_pair(
            KeyName=keypair["KeyName"],
        )

    def create_instance(self, keypair, security_group):
        print("Creating a source instance from {}".format(self.resource.ami))
        reservations = self.plan.client.run_instances(
            ImageId=self.resource.ami,
            InstanceType="m1.small",
            MaxCount=1,
            MinCount=1,
            KeyName=keypair['KeyName'],
            SecurityGroupIds=[security_group['GroupId']],
        )

        if len(reservations.get("Instances", [])) == 0:
            raise errors.Error("No instances were started")
        elif len(reservations["Instances"]) > 1:
            raise errors.Error("Somehow multiple instances were started!?")

        instance = reservations["Instances"][0]

        self.stack.callback(self.terminate_instance, instance)

        print("Waiting for instance {} to boot...".format(instance["InstanceId"]))
        self.plan.client.get_waiter("instance_running").wait(InstanceIds=[instance["InstanceId"]])

        return instance

    def terminate_instance(self, instance):
        print("Terminating instance")
        self.plan.client.terminate_instances(
            InstanceIds=[instance["InstanceId"]],
        )

        print("Waiting for instance to go away")
        self.plan.client.get_waiter("instance_terminated").wait(InstanceIds=[instance["InstanceId"]])

    def run(self):
        self.stack = ExitStack()
        with self.stack:
            keypair = self.create_keypair()
            security_group = self.create_security_group()
            instance = self.create_instance(keypair, security_group)

            print("Creating image")
            self.plan.client.create_image(
                Name=self.resource.name,
                InstanceId=instance['InstanceId'],
            )

            print("Waiting for image to become available")
            # self.plan.client.get_waiter("ami_available").wait(ImageId=[image["ImageId"]])
            for i in range(20):
                self.plan.object = self.plan.describe_object()
                if self.plan.object.get("State", "pending") == "available":
                    return
                time.sleep(15)
            raise errors.Error("Image didn't become available")

class Describe(SimpleDescribe, Plan):

    resource = Image
    service_name = 'ec2'
    describe_action = "describe_images"
    describe_list_key = "Images"
    key = 'ImageId'

    def get_describe_filters(self):
        return {"Filters": [{"Name": "name", "Values": [self.resource.name]}]}

class Apply(SimpleApply, Describe):

    create_action = "create_image"
    create_response = "not-that-useful"

    def create_object(self):
        return BuildInstance(self)

    def update_object(self):
        for change in super(Apply, self).update_object():
            yield change

class Destroy(SimpleDestroy, Describe):

    destroy_action = "deregister_image"
<commit_msg>ami -> source_ami<commit_after># Copyright 2014-2015 Isotoma Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import time
try:
    from contextlib import ExitStack
except ImportError:
    from contextlib2 import ExitStack

from touchdown.core.action import Action
from touchdown.core.resource import Resource
from touchdown.core.plan import Plan
from touchdown.core import argument, errors

from ..account import BaseAccount
from ..common import SimpleDescribe, SimpleApply, SimpleDestroy

class Image(Resource):

    resource_name = "image"

    name = argument.String(min=3, max=128, field="Name")
    description = argument.String(field="Description")

    source_ami = argument.String()

    resources = argument.List()

    #architecture = argument.String(field="Architecture", default="x86_64", choices=["x86_64", "i386"])
    #kernel = argument.String(field="KernelId")
    #ramdisk = argument.String(field="RamdiskId")
    #root_device_name = argument.String(field="RootDeviceName")
    #virtualization_type = argument.String(choices=["paravirtual", "hvm"], field="VirtualizationType")
    #sriov_net_support = argument.String(choices=["simple"], field="SriovNetSupport")
    #location = argument.String()
    #snapshot_id = argument.String()

    tags = argument.Dict()

    account = argument.Resource(BaseAccount)

class BuildInstance(Action):

    @property
    def description(self):
        yield "Build new AMI '{}' from '{}'".format(self.resource.name, self.resource.source_ami)

    def create_security_group(self):
        print("Creating temporary security group")
        security_group = self.plan.client.create_security_group(
            GroupName="temporary-security-group",
            Description="Temporary security group",
        )
        self.stack.callback(self.destroy_security_group, security_group)
        return security_group

    def destroy_security_group(self, security_group):
        print("Deleting temproary security group")
        self.plan.client.delete_security_group(
            GroupId=security_group["GroupId"],
        )

    def create_keypair(self):
        print("Creating temporary keypair")
        keypair = self.plan.client.create_key_pair(
            KeyName="temporary-key-pair",
        )
        self.stack.callback(self.destroy_keypair, keypair)
        return keypair

    def destroy_keypair(self, keypair):
        print("Deleting temproary keypair")
        self.plan.client.delete_key_pair(
            KeyName=keypair["KeyName"],
        )

    def create_instance(self, keypair, security_group):
        print("Creating a source instance from {}".format(self.resource.source_ami))
        reservations = self.plan.client.run_instances(
            ImageId=self.resource.source_ami,
            InstanceType="m1.small",
            MaxCount=1,
            MinCount=1,
            KeyName=keypair['KeyName'],
            SecurityGroupIds=[security_group['GroupId']],
        )

        if len(reservations.get("Instances", [])) == 0:
            raise errors.Error("No instances were started")
        elif len(reservations["Instances"]) > 1:
            raise errors.Error("Somehow multiple instances were started!?")

        instance = reservations["Instances"][0]

        self.stack.callback(self.terminate_instance, instance)

        print("Waiting for instance {} to boot...".format(instance["InstanceId"]))
        self.plan.client.get_waiter("instance_running").wait(InstanceIds=[instance["InstanceId"]])

        return instance

    def terminate_instance(self, instance):
        print("Terminating instance")
        self.plan.client.terminate_instances(
            InstanceIds=[instance["InstanceId"]],
        )

        print("Waiting for instance to go away")
        self.plan.client.get_waiter("instance_terminated").wait(InstanceIds=[instance["InstanceId"]])

    def run(self):
        self.stack = ExitStack()
        with self.stack:
            keypair = self.create_keypair()
            security_group = self.create_security_group()
            instance = self.create_instance(keypair, security_group)

            print("Creating image")
            self.plan.client.create_image(
                Name=self.resource.name,
                InstanceId=instance['InstanceId'],
            )

            print("Waiting for image to become available")
            # self.plan.client.get_waiter("ami_available").wait(ImageId=[image["ImageId"]])
            for i in range(20):
                self.plan.object = self.plan.describe_object()
                if self.plan.object.get("State", "pending") == "available":
                    return
                time.sleep(15)
            raise errors.Error("Image didn't become available")

class Describe(SimpleDescribe, Plan):

    resource = Image
    service_name = 'ec2'
    describe_action = "describe_images"
    describe_list_key = "Images"
    key = 'ImageId'

    def get_describe_filters(self):
        return {"Filters": [{"Name": "name", "Values": [self.resource.name]}]}

class Apply(SimpleApply, Describe):

    create_action = "create_image"
    create_response = "not-that-useful"

    def create_object(self):
        return BuildInstance(self)

    def update_object(self):
        for change in super(Apply, self).update_object():
            yield change

class Destroy(SimpleDestroy, Describe):

    destroy_action = "deregister_image"
<|endoftext|>