<commit_before>from argparse import ArgumentParser
from base64 import b64encode
from contextlib import contextmanager
import copy
from hashlib import sha512
from itertools import count
import json
import logging
import subprocess

import yaml

from jujupy import (
    EnvJujuClient,
    JESNotSupported,
    JujuData,
    SoftDeadlineExceeded,
)

__metaclass__ = type


def check_juju_output(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        if 'service' in result:
            raise AssertionError('Result contained service')
        return result
    return wrapper


class ControllerOperation(Exception):

    def __init__(self, operation):
        super(ControllerOperation, self).__init__(
            'Operation "{}" is only valid on controller models.'.format(
                operation))


def assert_juju_call(test_case, mock_method, client, expected_args,
                     call_index=None):
    if call_index is None:
        test_case.assertEqual(len(mock_method.mock_calls), 1)
        call_index = 0
    empty, args, kwargs = mock_method.mock_calls[call_index]
    test_case.assertEqual(args, (expected_args,))


class FakeControllerState:

    def __init__(self):
        self.name = 'name'
        self.state = 'not-bootstrapped'
        self.models = {}
        self.users = {
            'admin': {
                'state': '',
                'permission': 'write'
            }
        }
        self.shares = ['admin']

    def add_model(self, name):
        state = FakeEnvironmentState(self)
        state.name = name
        self.models[name] = state
        state.controller.state = 'created'
        return state

    def require_controller(self, operation, name):
        if name != self.controller_model.name:
            raise ControllerOperation(operation)

    def grant(self, username, permission):
        model_permissions = ['read', 'write', 'admin']
        if permission in model_permissions:
            permission = 'login'
        self.users[username]['access'] = permission

    def add_user_perms(self, username, permissions):
        self.users.update(
            {username: {'state': '', 'permission': permissions}})
        self.shares.append(username)

    def bootstrap(self, model_name, config, separate_controller):
        default_model = self.add_model(model_name)
        default_model.name = model_name
        if separate_controller:
            controller_model = default_model.controller.add_model('controller')
        else:
            controller_model = default_model
        self.controller_model = controller_model
        controller_model.state_servers.append(controller_model.add_machine())
        self.state = 'bootstrapped'
        default_model.model_config = copy.deepcopy(config)
        self.models[default_model.name] = default_model
        return default_model

    def destroy(self, kill=False):
        for model in self.models.values():
            model.kill_controller()
        self.models.clear()
        if kill:
            self.state = 'controller-killed'
        else:
            self.state = 'controller-destroyed'


class FakeEnvironmentState:
    """A Fake environment state that can be used by multiple FakeBackends."""

    def __init__(self, controller=None):
        self._clear()
        if controller is not None:
            self.controller = controller
        else:
            self.controller = FakeControllerState()

    def _clear(self):
        self.name = None
        self.machine_id_iter = count()
        self.state_servers = []
        self.services = {}
        self.machines = set()
        self.containers = {}
        self.relations = {}
        self.token = None
        self.exposed = set()
        self.machine_host_names = {}
        self.current_bundle = None
        self.model_config = None
        self.ssh_keys = []

    @property
    def state(self):
        return self.controller.state

    def add_machine(self, host_name=None, machine_id=None):
        if machine_id is None:
            machine_id = str(self.machine_id_iter.next())
        self.machines.add(machine_id)
        if host_name is None:
            host_name = '{}.example.com'.format(machine_id)
        self.machine_host_names[machine_id] = host_name
        return machine_id

    def add_ssh_machines(self, machines):
        for machine in machines:
            self.add_machine()

    def add_container(self, container_type, host=None, container_num=None):
        if host is None:
            host = self.add_machine()
        host_containers = self.containers.setdefault(host, set())
        if container_num is None:
            same_type_containers = [x for x in host_containers if
                                    container_type in x]
            container_num = len(same_type_containers)
        container_name = '{}/{}/{}'.format(host, container_type, container_num)
        host_containers.add(container_name)
        host_name = '{}.example.com'.format(container_name)
        self.machine_host_names[container_name] = host_name

    def remove_container(self, container_id):
        for containers in self.containers.values():
            containers.discard(container_id)

    def remove_machine(self, machine_id):
        self.machines.remove(machine_id)
        self.containers.pop(machine_id, None)

    def remove_state_server(self, machine_id):
        self.remove_machine(machine_id)
        self.state_servers.remove(machine_id)

    def destroy_environment(self):
        self._clear()
        self.controller.state = 'destroyed'
        return 0

    def kill_controller(self):
        self._clear()
        self.controller.state = 'controller-killed'

    def destroy_model(self):
        del self.controller.models[self.name]
        self._clear()
        self.controller.state = 'model-destroyed'

    def _fail_stderr(self, message, returncode=1, cmd='juju', stdout=''):
        exc = subprocess.CalledProcessError(returncode, cmd, stdout)
        exc.stderr = message
        raise exc

    def restore_backup(self):
        self.controller.require_controller('restore', self.name)
        if len(self.state_servers) > 0:
            return self._fail_stderr('Operation not permitted')
        self.state_servers.append(self.add_machine())

    def enable_ha(self):
        self.controller.require_controller('enable-ha', self.name)
        for n in range(2):<commit_msg>Switch fake from kill_controller to destroy_model, nix kill_controller.<commit_after>from argparse import ArgumentParser
from base64 import b64encode
from contextlib import contextmanager
import copy
from hashlib import sha512
from itertools import count
import json
import logging
import subprocess

import yaml

from jujupy import (
    EnvJujuClient,
    JESNotSupported,
    JujuData,
    SoftDeadlineExceeded,
)

__metaclass__ = type


def check_juju_output(func):
    def wrapper(*args, **kwargs):
        result = func(*args, **kwargs)
        if 'service' in result:
            raise AssertionError('Result contained service')
        return result
    return wrapper


class ControllerOperation(Exception):

    def __init__(self, operation):
        super(ControllerOperation, self).__init__(
            'Operation "{}" is only valid on controller models.'.format(
                operation))


def assert_juju_call(test_case, mock_method, client, expected_args,
                     call_index=None):
    if call_index is None:
        test_case.assertEqual(len(mock_method.mock_calls), 1)
        call_index = 0
    empty, args, kwargs = mock_method.mock_calls[call_index]
    test_case.assertEqual(args, (expected_args,))


class FakeControllerState:

    def __init__(self):
        self.name = 'name'
        self.state = 'not-bootstrapped'
        self.models = {}
        self.users = {
            'admin': {
                'state': '',
                'permission': 'write'
            }
        }
        self.shares = ['admin']

    def add_model(self, name):
        state = FakeEnvironmentState(self)
        state.name = name
        self.models[name] = state
        state.controller.state = 'created'
        return state

    def require_controller(self, operation, name):
        if name != self.controller_model.name:
            raise ControllerOperation(operation)

    def grant(self, username, permission):
        model_permissions = ['read', 'write', 'admin']
        if permission in model_permissions:
            permission = 'login'
        self.users[username]['access'] = permission

    def add_user_perms(self, username, permissions):
        self.users.update(
            {username: {'state': '', 'permission': permissions}})
        self.shares.append(username)

    def bootstrap(self, model_name, config, separate_controller):
        default_model = self.add_model(model_name)
        default_model.name = model_name
        if separate_controller:
            controller_model = default_model.controller.add_model('controller')
        else:
            controller_model = default_model
        self.controller_model = controller_model
        controller_model.state_servers.append(controller_model.add_machine())
        self.state = 'bootstrapped'
        default_model.model_config = copy.deepcopy(config)
        self.models[default_model.name] = default_model
        return default_model

    def destroy(self, kill=False):
        for model in self.models.values():
            model.destroy_model()
        self.models.clear()
        if kill:
            self.state = 'controller-killed'
        else:
            self.state = 'controller-destroyed'


class FakeEnvironmentState:
    """A Fake environment state that can be used by multiple FakeBackends."""

    def __init__(self, controller=None):
        self._clear()
        if controller is not None:
            self.controller = controller
        else:
            self.controller = FakeControllerState()

    def _clear(self):
        self.name = None
        self.machine_id_iter = count()
        self.state_servers = []
        self.services = {}
        self.machines = set()
        self.containers = {}
        self.relations = {}
        self.token = None
        self.exposed = set()
        self.machine_host_names = {}
        self.current_bundle = None
        self.model_config = None
        self.ssh_keys = []

    @property
    def state(self):
        return self.controller.state

    def add_machine(self, host_name=None, machine_id=None):
        if machine_id is None:
            machine_id = str(self.machine_id_iter.next())
        self.machines.add(machine_id)
        if host_name is None:
            host_name = '{}.example.com'.format(machine_id)
        self.machine_host_names[machine_id] = host_name
        return machine_id

    def add_ssh_machines(self, machines):
        for machine in machines:
            self.add_machine()

    def add_container(self, container_type, host=None, container_num=None):
        if host is None:
            host = self.add_machine()
        host_containers = self.containers.setdefault(host, set())
        if container_num is None:
            same_type_containers = [x for x in host_containers if
                                    container_type in x]
            container_num = len(same_type_containers)
        container_name = '{}/{}/{}'.format(host, container_type, container_num)
        host_containers.add(container_name)
        host_name = '{}.example.com'.format(container_name)
        self.machine_host_names[container_name] = host_name

    def remove_container(self, container_id):
        for containers in self.containers.values():
            containers.discard(container_id)

    def remove_machine(self, machine_id):
        self.machines.remove(machine_id)
        self.containers.pop(machine_id, None)

    def remove_state_server(self, machine_id):
        self.remove_machine(machine_id)
        self.state_servers.remove(machine_id)

    def destroy_environment(self):
        self._clear()
        self.controller.state = 'destroyed'
        return 0

    def destroy_model(self):
        del self.controller.models[self.name]
        self._clear()
        self.controller.state = 'model-destroyed'

    def _fail_stderr(self, message, returncode=1, cmd='juju', stdout=''):
        exc = subprocess.CalledProcessError(returncode, cmd, stdout)
        exc.stderr = message
        raise exc

    def restore_backup(self):
        self.controller.require_controller('restore', self.name)
        if len(self.state_servers) > 0:
            return self._fail_stderr('Operation not permitted')
        self.state_servers.append(self.add_machine())

    def enable_ha(self):
        self.controller.require_controller('enable-ha', self.name)
        for n in range(2):
{'fraction_non_alphanumeric': 0.09436905318982998, 'fraction_numerical': 0.0020277647792856027, 'mean_word_length': 2.513150684931507, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '285142', 'n_tokens_mistral': 3614, 'n_tokens_neox': 3522, 'n_words': 791}