<commit_before>from time import sleep

from django.contrib.auth.models import User
from django.core import mail
from tests.conftest import twilio_vcr

import pytest
from apostello.models import SmsOutbound

URI = '/config/first_run/'

@pytest.mark.django_db
@pytest.mark.slow
@pytest.mark.selenium
class TestFirstRun:
    def test_page_content(self, live_server, browser):
        browser.get(live_server + URI)
        assert URI in browser.current_url
        assert "Welcome" in browser.page_source
        assert "DJANGO_EMAIL_HOST" in browser.page_source
        assert "TWILIO_AUTH_TOKEN" in browser.page_source

    def disabled_test_email_form(self, live_server, browser, driver_wait_time):
        # TODO fix test - click does not work!
        browser.get(live_server + URI)
        assert URI in browser.current_url
        sleep(driver_wait_time)
        email_input_box = browser.find_elements_by_css_selector(
            '#send_test_email > div > form > div.fields > div.four.wide.field > input[type="email"]'
        )[0]
        email_input_box.clear()
        email_input_box.send_keys('test@example.com')

        body_input_box = browser.find_elements_by_css_selector(
            '#send_test_email > div > form > div.fields > div.twelve.wide.field > input[type="text"]'
        )[0]
        body_input_box.clear()
        body_input_box.send_keys('test message')

        submit_button = browser.find_elements_by_css_selector(
            '#send_test_email > div > form > button'
        )[0]
        import pdb
        pdb.set_trace()
        submit_button.click()

        sleep(driver_wait_time)
        assert len(mail.outbox) == 1
        assert 'test message' in mail.outbox[0].body

    @twilio_vcr
    def disabled_test_sms_form(
        self, live_server, browser, driver_wait_time, recipients
    ):
        # TODO fix test - click does not work!
        browser.get(live_server + URI)
        assert URI in browser.current_url
        sleep(driver_wait_time)
        to_input_box = browser.find_elements_by_css_selector(
            '#send_test_sms > div > form > div > div.four.wide.field > input[type="email"]'
        )[0]
        to_input_box.clear()
        to_input_box.send_keys(str(recipients['calvin'].number))

        body_input_box = browser.find_elements_by_css_selector(
            '#send_test_sms > div > form > div > div.twelve.wide.field > input[type="text"]'
        )[0]
        body_input_box.clear()
        body_input_box.send_keys('test')

        submit_button = browser.find_elements_by_css_selector(
            '#send_test_sms > div > form > button'
        )[0]
        submit_button.click()

        sleep(driver_wait_time)
        assert SmsOutbound.objects.count() == 1

    def disabled_test_user_form(self, live_server, browser, driver_wait_time):
        # TODO fix test - click does not work!
        browser.get(live_server + URI)
        assert URI in browser.current_url
        sleep(driver_wait_time)
        email_input_box = browser.find_elements_by_css_selector(
            '#create_admin_user > div > form > div > div.eight.wide.field > input[type="email"]'
        )[0]
        email_input_box.clear()
        email_input_box.send_keys('test@example.com')

        pass1_input_box = browser.find_elements_by_css_selector(
            '#create_admin_user > div > form > div > div:nth-child(2) > input[type="password"]'
        )[0]
        pass1_input_box.clear()
        pass1_input_box.send_keys('password')

        pass2_input_box = browser.find_elements_by_css_selector(
            '#create_admin_user > div > form > div > div:nth-child(3) > input[type="password"]'
        )[0]
        pass2_input_box.clear()
        pass2_input_box.send_keys('password')

        submit_button = browser.find_elements_by_css_selector(
            '#create_admin_user > div > form > button'
        )[0]
        submit_button.click()

        sleep(driver_wait_time)
        assert User.objects.count() == 1
        browser.get(live_server + URI)
        assert URI not in browser.current_url
<commit_msg>fix first run tests<commit_after>from time import sleep

from django.contrib.auth.models import User
from django.core import mail
from tests.conftest import twilio_vcr

import pytest
from apostello.models import SmsOutbound

URI = '/config/first_run/'

@pytest.mark.django_db
@pytest.mark.slow
@pytest.mark.selenium
class TestFirstRun:
    def test_page_content(self, live_server, browser):
        browser.get(live_server + URI)
        assert URI in browser.current_url
        assert "Welcome" in browser.page_source
        assert "DJANGO_EMAIL_HOST" in browser.page_source
        assert "TWILIO_AUTH_TOKEN" in browser.page_source

    def test_email_form(self, live_server, browser, driver_wait_time):
        browser.get(live_server + URI)
        assert URI in browser.current_url
        sleep(driver_wait_time)
        email_input_box = browser.find_elements_by_css_selector(
            '#send_test_email > div > form > div.fields > div.four.wide.field > input[type="email"]'
        )[0]
        email_input_box.clear()
        email_input_box.send_keys('test@example.com')

        body_input_box = browser.find_elements_by_css_selector(
            '#send_test_email > div > form > div.fields > div.twelve.wide.field > input[type="text"]'
        )[0]
        body_input_box.clear()
        body_input_box.send_keys('test message')

        submit_button = browser.find_elements_by_css_selector(
            '#send_test_email > div > form > button'
        )[0]
        submit_button.click()

        sleep(driver_wait_time)
        assert len(mail.outbox) == 1
        assert 'test message' in mail.outbox[0].body

    # @twilio_vcr
    def test_sms_form(
        self, live_server, browser, driver_wait_time, recipients
    ):
        browser.get(live_server + URI)
        assert URI in browser.current_url
        sleep(driver_wait_time)
        to_input_box = browser.find_elements_by_css_selector(
            '#send_test_sms > div > form > div > div.four.wide.field > input[type="text"]'
        )[0]
        to_input_box.clear()
        to_input_box.send_keys(str(recipients['calvin'].number))

        body_input_box = browser.find_elements_by_css_selector(
            '#send_test_sms > div > form > div > div.twelve.wide.field > input[type="text"]'
        )[0]
        body_input_box.clear()
        body_input_box.send_keys('test')

        submit_button = browser.find_elements_by_css_selector(
            '#send_test_sms > div > form > button'
        )[0]
        submit_button.click()

        sleep(driver_wait_time)
        assert 'AC00000000000000000000000000000000' in browser.page_source
        assert 'Twilio returned the following information:' in browser.page_source

    def test_user_form(self, live_server, browser, driver_wait_time):
        browser.get(live_server + URI)
        assert URI in browser.current_url
        sleep(driver_wait_time)
        email_input_box = browser.find_elements_by_css_selector(
            '#create_admin_user > div > form > div > div.eight.wide.field > input[type="email"]'
        )[0]
        email_input_box.clear()
        email_input_box.send_keys('test@example.com')

        pass1_input_box = browser.find_elements_by_css_selector(
            '#create_admin_user > div > form > div > div:nth-child(2) > input[type="password"]'
        )[0]
        pass1_input_box.clear()
        pass1_input_box.send_keys('password')

        pass2_input_box = browser.find_elements_by_css_selector(
            '#create_admin_user > div > form > div > div:nth-child(3) > input[type="password"]'
        )[0]
        pass2_input_box.clear()
        pass2_input_box.send_keys('password')

        submit_button = browser.find_elements_by_css_selector(
            '#create_admin_user > div > form > button'
        )[0]
        submit_button.click()

        sleep(driver_wait_time)
        assert User.objects.count() == 1
        browser.get(live_server + URI)
        assert URI not in browser.current_url
<|endoftext|>


<commit_before><commit_msg>planetoid model description<commit_after><|endoftext|>


<commit_before>
"""Resource manager for apache-solr-jetty-server 3.5
"""

# Common stdlib imports
import sys
import os
import os.path
## import commands

# fix path if necessary (if running from source or running as test)
try:
    import engage.utils
except:
    sys.exc_clear()
    dir_to_add_to_python_path = os.path.abspath((os.path.join(os.path.dirname(__file__), "../../../..")))
    sys.path.append(dir_to_add_to_python_path)

import engage.drivers.service_manager as service_manager
import engage.drivers.utils
# Drivers compose *actions* to implement their methods.
from engage.drivers.action import *

# setup errors
from engage.utils.user_error import UserError, EngageErrInf
import gettext
_ = gettext.gettext

errors = { }

def define_error(error_code, msg):
    global errors
    error_info = EngageErrInf(__name__, error_code, msg)
    errors[error_info.error_code] = error_info

# error codes
# FILL IN
ERR_TBD = 0

define_error(ERR_TBD,
             _("Replace this with your error codes"))

# setup logging
from engage.utils.log_setup import setup_engage_logger
logger = setup_engage_logger(__name__)

# this is used by the package manager to locate the packages.json
# file associated with the driver
def get_packages_filename():
    return engage.drivers.utils.get_packages_filename(__file__)

def make_context(resource_json, sudo_password_fn, dry_run=False):
    """Create a Context object (defined in engage.utils.action). This contains
    the resource's metadata in ctx.props, references to the logger and sudo
    password function, and various helper functions. The context object is used
    by individual actions.

    If your resource does not need the sudo password, you can just pass in
    None for sudo_password_fn.
    """
    ctx = Context(resource_json, logger, __file__,
                  sudo_password_fn=sudo_password_fn,
                  dry_run=dry_run)
    ctx.check_port('input_ports.host',
                  genforma_home=unicode,
                  log_directory=unicode)
    ctx.check_port('input_ports.schema_file',
                  file_path=unicode)
    ctx.check_port('input_ports.jvm',
                  java_exe=unicode)
    ctx.check_port('output_ports.solr',
                  home=unicode)

    # add any extra computed properties here using the ctx.add() method.
    home = ctx.props.output_ports.solr.home
    example = os.path.join(home, "example")
    ctx.add('startup_jar_file',
            os.path.join(example, "start.jar"))
    ctx.add('schema_file_target',
            os.path.join(example, 'solr/conf/schema.xml'))
    ctx.add('pid_file',
            os.path.join(home, "apache_solr.pid"))
    ctx.add('log_file',
            os.path.join(ctx.props.input_ports.host.log_directory,
                         "apache_solr.log"))
    return ctx

class Manager(service_manager.Manager):
    REQUIRES_ROOT_ACCESS = False
    def __init__(self, metadata, dry_run=False):
        package_name = "%s %s" % (metadata.key["name"],
                                  metadata.key["version"])
        service_manager.Manager.__init__(self, metadata, package_name)
        self.ctx = make_context(metadata.to_json(),
                                None,
                                dry_run=dry_run)

    def validate_pre_install(self):
        p = self.ctx.props
        self.ctx.r(check_installable_to_dir,
                   p.output_ports.solr.home)

    def is_installed(self):
        return os.path.exists(self.ctx.props.output_ports.solr.home)

    def install(self, package):
        p = self.ctx.props
        self.ctx.r(extract_package_as_dir, package,
                   p.output_ports.solr.home)
        self.ctx.r(wrap_action(shutil.copy),
                   p.input_ports.schema_file.file_path,
                   p.schema_file_target)

    def validate_post_install(self):
        p = self.ctx.props
        self.ctx.r(check_dir_exists,  p.output_ports.solr.home)

    def start(self):
        p = self.ctx.props
        command_exe = p.input_ports.jvm.java_exe
        self.ctx.r(start_server,
                   [command_exe, "-jar", p.startup_jar_file],
                   p.log_file,
                   p.pid_file,
                   cwd=os.path.dirname(p.startup_jar_file))

    def is_running(self):
        p = self.ctx.props
        return self.ctx.rv(get_server_status,
                           p.pid_file) != None

    def stop(self):
        p = self.ctx.props
        self.ctx.r(stop_server, p.pid_file)

    def get_pid_file_path(self):
        return self.ctx.props.pid_file
<commit_msg>Make sure solr is really accepting requests before returning from start()<commit_after>
"""Resource manager for apache-solr-jetty-server 3.5
"""

# Common stdlib imports
import sys
import os
import os.path
import urllib
from contextlib import closing

# fix path if necessary (if running from source or running as test)
try:
    import engage.utils
except:
    sys.exc_clear()
    dir_to_add_to_python_path = os.path.abspath((os.path.join(os.path.dirname(__file__), "../../../..")))
    sys.path.append(dir_to_add_to_python_path)

import engage.drivers.service_manager as service_manager
import engage.drivers.utils

# Drivers compose *actions* to implement their methods.
from engage.drivers.action import *

# setup errors
from engage.utils.user_error import UserError, EngageErrInf
import gettext
_ = gettext.gettext

errors = { }

def define_error(error_code, msg):
    global errors
    error_info = EngageErrInf(__name__, error_code, msg)
    errors[error_info.error_code] = error_info

# error codes
# FILL IN
ERR_TBD = 0

define_error(ERR_TBD,
             _("Replace this with your error codes"))

# setup logging
from engage.utils.log_setup import setup_engage_logger
logger = setup_engage_logger(__name__)

# this is used by the package manager to locate the packages.json
# file associated with the driver
def get_packages_filename():
    return engage.drivers.utils.get_packages_filename(__file__)

def make_context(resource_json, sudo_password_fn, dry_run=False):
    """Create a Context object (defined in engage.utils.action). This contains
    the resource's metadata in ctx.props, references to the logger and sudo
    password function, and various helper functions. The context object is used
    by individual actions.

    If your resource does not need the sudo password, you can just pass in
    None for sudo_password_fn.
    """
    ctx = Context(resource_json, logger, __file__,
                  sudo_password_fn=sudo_password_fn,
                  dry_run=dry_run)
    ctx.check_port('input_ports.host',
                  genforma_home=unicode,
                  log_directory=unicode)
    ctx.check_port('input_ports.schema_file',
                  file_path=unicode)
    ctx.check_port('input_ports.jvm',
                  java_exe=unicode)
    ctx.check_port('output_ports.solr',
                  home=unicode)

    # add any extra computed properties here using the ctx.add() method.
    home = ctx.props.output_ports.solr.home
    example = os.path.join(home, "example")
    ctx.add('startup_jar_file',
            os.path.join(example, "start.jar"))
    ctx.add('schema_file_target',
            os.path.join(example, 'solr/conf/schema.xml'))
    ctx.add('pid_file',
            os.path.join(home, "apache_solr.pid"))
    ctx.add('log_file',
            os.path.join(ctx.props.input_ports.host.log_directory,
                         "apache_solr.log"))
    return ctx

@make_value_action
def is_solr_alive(self, solr_url='http://localhost:8983/solr/select'):
    try:
        with closing(urllib.urlopen(solr_url)) as aiu:
            aiu.read()
        return True
    except IOError:
        return False

class Manager(service_manager.Manager):
    REQUIRES_ROOT_ACCESS = False
    def __init__(self, metadata, dry_run=False):
        package_name = "%s %s" % (metadata.key["name"],
                                  metadata.key["version"])
        service_manager.Manager.__init__(self, metadata, package_name)
        self.ctx = make_context(metadata.to_json(),
                                None,
                                dry_run=dry_run)

    def validate_pre_install(self):
        p = self.ctx.props
        self.ctx.r(check_installable_to_dir,
                   p.output_ports.solr.home)

    def is_installed(self):
        return os.path.exists(self.ctx.props.output_ports.solr.home)

    def install(self, package):
        p = self.ctx.props
        self.ctx.r(extract_package_as_dir, package,
                   p.output_ports.solr.home)
        self.ctx.r(wrap_action(shutil.copy),
                   p.input_ports.schema_file.file_path,
                   p.schema_file_target)

    def validate_post_install(self):
        p = self.ctx.props
        self.ctx.r(check_dir_exists,  p.output_ports.solr.home)

    def start(self):
        p = self.ctx.props
        command_exe = p.input_ports.jvm.java_exe
        self.ctx.r(start_server,
                   [command_exe, "-jar", p.startup_jar_file],
                   p.log_file,
                   p.pid_file,
                   cwd=os.path.dirname(p.startup_jar_file))
        self.ctx.check_poll(10, 20, lambda r: r, is_solr_alive)

    def is_running(self):
        p = self.ctx.props
        return self.ctx.rv(get_server_status,
                           p.pid_file) != None

    def stop(self):
        p = self.ctx.props
        self.ctx.r(stop_server, p.pid_file)

    def get_pid_file_path(self):
        return self.ctx.props.pid_file
<|endoftext|>


<commit_before>import json

from django.conf import settings
from django.core.urlresolvers import reverse
from django.core.urlresolvers import reverse_lazy
from django.http import HttpResponse  # noqa
from django.utils.translation import ugettext_lazy as _
from django.views.generic import View  # noqa

from horizon import exceptions
from horizon import views

from openstack_dashboard import api
from openstack_dashboard.usage import quotas

from openstack_dashboard.dashboards.project.network_topology.instances \
    import tables as instances_tables
from openstack_dashboard.dashboards.project.network_topology.ports \
    import tables as ports_tables
from openstack_dashboard.dashboards.project.network_topology.routers \
    import tables as routers_tables

from openstack_dashboard.dashboards.project.instances import\
    console as i_console
from openstack_dashboard.dashboards.project.instances import\
    views as i_views
from openstack_dashboard.dashboards.project.instances.workflows import\
    create_instance as i_workflows
from openstack_dashboard.dashboards.project.networks import\
    views as n_views
from openstack_dashboard.dashboards.project.networks import\
    workflows as n_workflows
from openstack_dashboard.dashboards.project.routers import\
    views as r_views

class IndexView(views.HorizonTemplateView):
    template_name = 'vdc/infrastructure/index.html'
    page_title = _("Infrastructure")

    def get_context_data(self, **kwargs):
        context = super(IndexView, self).get_context_data(**kwargs)
        return context

class JSONView(View):

    @property
    def add_resource_url(self, view, resources):
        tenant_id = self.request.user.tenant_id
        for resource in resources:
            if (resource.get('tenant_id')
                    and tenant_id != resource.get('tenant_id')):
                continue
            resource['url'] = reverse(view, None, [str(resource['id'])])

    def _get_servers(self, request):
        # Get nova data
        try:
            servers, more = api.nova.server_list(request)
        except Exception:
            servers = []
        data = []
        console_type = getattr(settings, 'CONSOLE_TYPE', 'AUTO')
        # lowercase of the keys will be used at the end of the console URL.
        for server in servers:
            try:
                console = i_console.get_console(
                    request, console_type, server)[0].lower()
            except exceptions.NotAvailable:
                console = None

            server_data = {'name': server.name,
                           'status': server.status,
                           'task': getattr(server, 'OS-EXT-STS:task_state'),
                           'id': server.id}
            if console:
                server_data['console'] = console
            data.append(server_data)
        #self.add_resource_url('horizon:project:instances:detail', data)
        return data

    def get(self, request, *args, **kwargs):
        data = {'servers': self._get_servers(request)}
        json_string = json.dumps(data, ensure_ascii=False)
        return HttpResponse(json_string, content_type='text/json')
<commit_msg>add get_network in view<commit_after>import json

from django.conf import settings
from django.core.urlresolvers import reverse
from django.core.urlresolvers import reverse_lazy
from django.http import HttpResponse  # noqa
from django.utils.translation import ugettext_lazy as _
from django.views.generic import View  # noqa

from horizon import exceptions
from horizon import views

from openstack_dashboard import api
from openstack_dashboard.usage import quotas

from openstack_dashboard.dashboards.project.network_topology.instances \
    import tables as instances_tables
from openstack_dashboard.dashboards.project.network_topology.ports \
    import tables as ports_tables
from openstack_dashboard.dashboards.project.network_topology.routers \
    import tables as routers_tables

from openstack_dashboard.dashboards.project.instances import\
    console as i_console
from openstack_dashboard.dashboards.project.instances import\
    views as i_views
from openstack_dashboard.dashboards.project.instances.workflows import\
    create_instance as i_workflows
from openstack_dashboard.dashboards.project.networks import\
    views as n_views
from openstack_dashboard.dashboards.project.networks import\
    workflows as n_workflows
from openstack_dashboard.dashboards.project.routers import\
    views as r_views

class IndexView(views.HorizonTemplateView):
    template_name = 'vdc/infrastructure/index.html'
    page_title = _("Infrastructure")

    def get_context_data(self, **kwargs):
        context = super(IndexView, self).get_context_data(**kwargs)
        return context

class JSONView(View):

    @property
    def add_resource_url(self, view, resources):
        tenant_id = self.request.user.tenant_id
        for resource in resources:
            if (resource.get('tenant_id')
                    and tenant_id != resource.get('tenant_id')):
                continue
            resource['url'] = reverse(view, None, [str(resource['id'])])

    def _get_servers(self, request):
        # Get nova data
        try:
            servers, more = api.nova.server_list(request)
        except Exception:
            servers = []
        data = []
        console_type = getattr(settings, 'CONSOLE_TYPE', 'AUTO')
        # lowercase of the keys will be used at the end of the console URL.
        for server in servers:
            try:
                console = i_console.get_console(
                    request, console_type, server)[0].lower()
            except exceptions.NotAvailable:
                console = None

            server_data = {'name': server.name,
                           'status': server.status,
                           'task': getattr(server, 'OS-EXT-STS:task_state'),
                           'id': server.id}
            if console:
                server_data['console'] = console
            data.append(server_data)
        #self.add_resource_url('horizon:project:instances:detail', data)
        return data

    def _get_networks(self, request):
        # Get neutron data
        # if we didn't specify tenant_id, all networks shown as admin user.
        # so it is need to specify the networks. However there is no need to
        # specify tenant_id for subnet. The subnet which belongs to the public
        # network is needed to draw subnet information on public network.
        try:
            neutron_networks = api.neutron.network_list_for_tenant(
                request,
                request.user.tenant_id)
        except Exception:
            neutron_networks = []
        networks = [{'name': network.name,
                     'id': network.id,
                     'subnets': [{'cidr': subnet.cidr}
                                 for subnet in network.subnets],
                     'router:external': network['router:external']}
                    for network in neutron_networks]
        self.add_resource_url('horizon:project:networks:detail',
                              networks)

        # Add public networks to the networks list
        if self.is_router_enabled:
            try:
                neutron_public_networks = api.neutron.network_list(
                    request,
                    **{'router:external': True})
            except Exception:
                neutron_public_networks = []
            my_network_ids = [net['id'] for net in networks]
            for publicnet in neutron_public_networks:
                if publicnet.id in my_network_ids:
                    continue
                try:
                    subnets = [{'cidr': subnet.cidr}
                               for subnet in publicnet.subnets]
                except Exception:
                    subnets = []
                networks.append({
                    'name': publicnet.name,
                    'id': publicnet.id,
                    'subnets': subnets,
                    'router:external': publicnet['router:external']})

        return sorted(networks,
                      key=lambda x: x.get('router:external'),
                      reverse=True)

    def get(self, request, *args, **kwargs):
        data = {'servers': self._get_servers(request)}
        json_string = json.dumps(data, ensure_ascii=False)
        return HttpResponse(json_string, content_type='text/json')
<|endoftext|>


<commit_before>from pytest import approx
import numpy as np
import pandas as pd
import motmetrics as mm
import os
import pytest

def test_events():
    acc = mm.MOTAccumulator()

    # All FP
    acc.update([], ['a', 'b'], [], frameid=0)
    # All miss
    acc.update([1, 2], [], [], frameid=1)
    # Match
    acc.update([1, 2], ['a', 'b'], [[1, 0.5], [0.3, 1]], frameid=2)
    # Switch
    acc.update([1, 2], ['a', 'b'], [[0.2, np.nan], [np.nan, 0.1]], frameid=3)
    # Match. Better new match is available but should prefer history
    acc.update([1, 2], ['a', 'b'], [[5, 1], [1, 5]], frameid=4)
    # No data
    acc.update([], [], [], frameid=5)

    expect = mm.MOTAccumulator.new_event_dataframe()
    expect.loc[(0, 0), :] = ['FP', np.nan, 'a', np.nan]
    expect.loc[(0, 1), :] = ['FP', np.nan, 'b', np.nan]
    expect.loc[(1, 0), :] = ['MISS', 1, np.nan, np.nan]
    expect.loc[(1, 1), :] = ['MISS', 2, np.nan, np.nan]
    expect.loc[(2, 0), :] = ['MATCH', 1, 'b', 0.5]
    expect.loc[(2, 1), :] = ['MATCH', 2, 'a', 0.3]
    expect.loc[(3, 0), :] = ['SWITCH', 1, 'a', 0.2]
    expect.loc[(3, 1), :] = ['SWITCH', 2, 'b', 0.1]
    expect.loc[(4, 0), :] = ['MATCH', 1, 'a', 5.]
    expect.loc[(4, 1), :] = ['MATCH', 2, 'b', 5.]
    # frame 5 generates no events

    assert pd.DataFrame.equals(acc.events, expect)

    mh = mm.metrics.create()
    metr = mh.compute(acc, metrics=['motp', 'mota'], return_dataframe=False)
    assert metr['motp'] == approx(11.1 / 6)
    assert metr['mota'] == approx(1. - (2 + 2 + 2) / 8)

def test_max_switch_time():
     acc = mm.MOTAccumulator(max_switch_time=1)
     df = acc.update([1, 2], ['a', 'b'], [[1, 0.5], [0.3, 1]], frameid=1) # 1->a, 2->b
     df = acc.update([1, 2], ['a', 'b'], [[0.5, np.nan], [np.nan, 0.5]], frameid=2) # 1->b, 2->a
     assert (df.Type == 'SWITCH').all()

     acc = mm.MOTAccumulator(max_switch_time=1)
     df = acc.update([1, 2], ['a', 'b'], [[1, 0.5], [0.3, 1]], frameid=1) # 1->a, 2->b
     df = acc.update([1, 2], ['a', 'b'], [[0.5, np.nan], [np.nan, 0.5]], frameid=5) # Later frame 1->b, 2->a
     assert (df.Type == 'MATCH').all()

def test_auto_id():
    acc = mm.MOTAccumulator(auto_id=True)
    acc.update([1, 2, 3, 4], [], [])
    acc.update([1, 2, 3, 4], [], [])
    assert acc.events.index.levels[0][-1] == 1
    acc.update([1, 2, 3, 4], [], [])
    assert acc.events.index.levels[0][-1] == 2

    with pytest.raises(AssertionError):
        acc.update([1, 2, 3, 4], [], [], frameid=5)

    acc = mm.MOTAccumulator(auto_id=False)
    with pytest.raises(AssertionError):
        acc.update([1, 2, 3, 4], [], [])

def test_correct_average():
    # Tests what is being depicted in figure 3 of 'Evaluating MOT Performance'
    acc = mm.MOTAccumulator(auto_id=True)

    # No track
    acc.update([1, 2, 3, 4], [], [])
    acc.update([1, 2, 3, 4], [], [])
    acc.update([1, 2, 3, 4], [], [])
    acc.update([1, 2, 3, 4], [], [])

    # Track single
    acc.update([4], [4], [0])
    acc.update([4], [4], [0])
    acc.update([4], [4], [0])
    acc.update([4], [4], [0])

    mh = mm.metrics.create()
    metr = mh.compute(acc, metrics='mota', return_dataframe=False)
    assert metr['mota'] == approx(0.2)<commit_msg>Adapted to raw events<commit_after>from pytest import approx
import numpy as np
import pandas as pd
import motmetrics as mm
import os
import pytest

def test_events():
    acc = mm.MOTAccumulator()

    # All FP
    acc.update([], ['a', 'b'], [], frameid=0)
    # All miss
    acc.update([1, 2], [], [], frameid=1)
    # Match
    acc.update([1, 2], ['a', 'b'], [[1, 0.5], [0.3, 1]], frameid=2)
    # Switch
    acc.update([1, 2], ['a', 'b'], [[0.2, np.nan], [np.nan, 0.1]], frameid=3)
    # Match. Better new match is available but should prefer history
    acc.update([1, 2], ['a', 'b'], [[5, 1], [1, 5]], frameid=4)
    # No data
    acc.update([], [], [], frameid=5)

    expect = mm.MOTAccumulator.new_event_dataframe()
    expect.loc[(0, 0), :] = ['FP', np.nan, 'a', np.nan]
    expect.loc[(0, 1), :] = ['FP', np.nan, 'b', np.nan]
    expect.loc[(1, 0), :] = ['MISS', 1, np.nan, np.nan]
    expect.loc[(1, 1), :] = ['MISS', 2, np.nan, np.nan]

    expect.loc[(2, 0), :] = ['RAW', 1, 'a', 1.0]
    expect.loc[(2, 1), :] = ['RAW', 1, 'b', 0.5]
    expect.loc[(2, 2), :] = ['RAW', 2, 'a', 0.3]
    expect.loc[(2, 3), :] = ['RAW', 2, 'b', 1.0]
    expect.loc[(2, 4), :] = ['MATCH', 1, 'b', 0.5]
    expect.loc[(2, 5), :] = ['MATCH', 2, 'a', 0.3]

    expect.loc[(3, 0), :] = ['RAW', 1, 'a', 0.2]
    expect.loc[(3, 1), :] = ['RAW', 2, 'b', 0.1]
    expect.loc[(3, 2), :] = ['SWITCH', 1, 'a', 0.2]
    expect.loc[(3, 3), :] = ['SWITCH', 2, 'b', 0.1]

    expect.loc[(4, 0), :] = ['RAW', 1, 'a', 5.]
    expect.loc[(4, 1), :] = ['RAW', 1, 'b', 1.]
    expect.loc[(4, 2), :] = ['RAW', 2, 'a', 1.]
    expect.loc[(4, 3), :] = ['RAW', 2, 'b', 5.]
    expect.loc[(4, 4), :] = ['MATCH', 1, 'a', 5.]
    expect.loc[(4, 5), :] = ['MATCH', 2, 'b', 5.]
    # frame 5 generates no events

    assert pd.DataFrame.equals(acc.events, expect)

def test_max_switch_time():
     acc = mm.MOTAccumulator(max_switch_time=1)
     df = acc.update([1, 2], ['a', 'b'], [[1, 0.5], [0.3, 1]], frameid=1) # 1->a, 2->b
     df = acc.update([1, 2], ['a', 'b'], [[0.5, np.nan], [np.nan, 0.5]], frameid=2) # 1->b, 2->a

     df = df[df.Type != 'RAW']
     assert (df.Type == 'SWITCH').all()

     acc = mm.MOTAccumulator(max_switch_time=1)
     df = acc.update([1, 2], ['a', 'b'], [[1, 0.5], [0.3, 1]], frameid=1) # 1->a, 2->b
     df = acc.update([1, 2], ['a', 'b'], [[0.5, np.nan], [np.nan, 0.5]], frameid=5) # Later frame 1->b, 2->a
     df = df[df.Type != 'RAW']
     assert (df.Type == 'MATCH').all()

def test_auto_id():
    acc = mm.MOTAccumulator(auto_id=True)
    acc.update([1, 2, 3, 4], [], [])
    acc.update([1, 2, 3, 4], [], [])
    assert acc.events.index.levels[0][-1] == 1
    acc.update([1, 2, 3, 4], [], [])
    assert acc.events.index.levels[0][-1] == 2

    with pytest.raises(AssertionError):
        acc.update([1, 2, 3, 4], [], [], frameid=5)

    acc = mm.MOTAccumulator(auto_id=False)
    with pytest.raises(AssertionError):
        acc.update([1, 2, 3, 4], [], [])<|endoftext|>


<commit_before>import json

from django.utils import timezone
from django.conf import settings
from django.core.exceptions import ValidationError
from django.db import models
from django.utils.translation import ugettext_lazy

from devilry.devilry_comment import models as comment_models
from devilry.apps.core.models import assignment_group

class AbstractGroupCommentQuerySet(models.QuerySet):
    """
    Base class for QuerySets for :class:`.AbstractGroupComment`.
    """
    def exclude_private_comments_from_other_users(self, user):
        return self.exclude(
            models.Q(visibility=AbstractGroupComment.VISIBILITY_PRIVATE) &
            ~models.Q(user=user)
        )

    def exclude_is_part_of_grading_feedbackset_unpublished(self):
        return self.exclude(
            part_of_grading=True,
            feedback_set__grading_published_datetime__isnull=True
        )

class AbstractGroupComment(comment_models.Comment):
    """
    The abstract superclass of all comments related to a delivery and feedback
    """

    #: The related feedbackset. See :class:`.FeedbackSet`.
    feedback_set = models.ForeignKey('FeedbackSet')

    #: If this is ``True``, the comment is published when the feedbackset
    #: is published. This means that this comment is part of the feedback/grading
    #: from the examiner. The same :obj:`~.AbstractGroupComment.visibility`
    #: rules apply no matter if this is ``True`` or ``False``,
    #: this only controls when the comment is published.
    part_of_grading = models.BooleanField(default=False)

    #: Comment only visible for :obj:`~devilry_comment.models.Comment.user` that created comment.
    #: Choice for :obj:`~.AbstractGroupComment.visibility`.
    VISIBILITY_PRIVATE = 'private'

    #: Comment should only be visible to examiners and admins that has
    #: access to the :obj:`~.AbstractGroupComment.feedback_set`.
    #: Choice for :obj:`~.AbstractGroupComment.visibility`.
    VISIBILITY_VISIBLE_TO_EXAMINER_AND_ADMINS = 'visible-to-examiner-and-admins'

    #: Comment should be visible to everyone that has
    #: access to the :obj:`~.AbstractGroupComment.feedback_set`.
    #: Choice for :obj:`~.AbstractGroupComment.visibility`.
    VISIBILITY_VISIBLE_TO_EVERYONE = 'visible-to-everyone'

    #: Choice list.
    #: Choices for :obj:`~.AbstractGroupComment.visibility`.
    VISIBILITY_CHOICES = [
        (VISIBILITY_PRIVATE, 'Private'),
        (VISIBILITY_VISIBLE_TO_EXAMINER_AND_ADMINS, 'Visible to examiners and admins'),
        (VISIBILITY_VISIBLE_TO_EVERYONE, 'Visible to everyone'),
    ]

    #: Sets the visibility choise of the comment.
    #: Defaults to :obj:`~.AbstractGroupComment.VISIBILITY_VISIBLE_TO_EVERYONE`.
    visibility = models.CharField(
        max_length=50,
        db_index=True,
        choices=VISIBILITY_CHOICES,
        default=VISIBILITY_VISIBLE_TO_EVERYONE,
    )

    class Meta:
        abstract = True

    def get_published_datetime(self):
        """
        Get the publishing datetime of the comment. Publishing datetime is
        the publishing time of the FeedbackSet if the comment has
        :obj:`~devilry.devilry_group.models.AbstractGroupComment.part_of_grading`
        set to True.

        Returns: Datetime.

        """
        return self.feedback_set.grading_published_datetime \
            if self.part_of_grading \
            else self.published_datetime

    def publish_draft(self, time):
        self.published_datetime = time + timezone.timedelta(milliseconds=1)
        self.full_clean()
        self.save()

class FeedbackSet(models.Model):
    """
    All comments that are given for a specific deadline (delivery and feedback) are
    linked to a feedback-set.

    If the comment has `instant_publish=True` it will be published instantly, otherwise the comments will only be
    visible once the feedbackset is published.
    All student-comments will be `instant_publish=True`, and the same applies to comments made by examiners that
    are not a part of feedback.
    """

    #: The AssignmentGroup that owns this feedbackset.
    group = models.ForeignKey(assignment_group.AssignmentGroup)

    #: Is the last feedbackset for :obj:`~.FeedbackSet.group`? Must be None or True.
    is_last_in_group = models.NullBooleanField(default=True)

    #: This means the feedbackset is basically the first feedbackset.
    #: Choice for :obj:`~.FeedbackSet.feedbackset_type`.
    FEEDBACKSET_TYPE_FIRST_ATTEMPT = 'first_attempt'

    #: Is not the first feedbackset, but a new attempt.
    #: Choice for :obj:`~.FeedbackSet.feedbackset_type`
    FEEDBACKSET_TYPE_NEW_ATTEMPT = 'new_attempt'

    #: Something went wrong on grading, with this option, a new
    #: deadline should not be given to student. Student should just
    #: get notified that a new feedback was given.
    #: Choice for :obj:`~.FeedbackSet.feedbackset_type`.
    FEEDBACKSET_TYPE_RE_EDIT = 're_edit'

    #: Grading status choices for :obj:`~.FeedbackSet.feedbackset_type`.
    FEEDBACKSET_TYPE_CHOICES = [
        (FEEDBACKSET_TYPE_FIRST_ATTEMPT, 'first attempt'),
        (FEEDBACKSET_TYPE_NEW_ATTEMPT,'new attempt'),
        (FEEDBACKSET_TYPE_RE_EDIT, 're edit'),
    ]

    #: Sets the type of the feedbackset.
    #: Defaults to :obj:`~.FeedbackSet.FEEDBACKSET_TYPE_FIRST_ATTEMPT`.
    feedbackset_type = models.CharField(
        max_length=50,
        db_index=True,
        choices=FEEDBACKSET_TYPE_CHOICES,
        default=FEEDBACKSET_TYPE_FIRST_ATTEMPT,
    )

    #: The User that created the feedbackset. Only used as metadata
    #: for superusers (for debugging).
    created_by = models.ForeignKey(settings.AUTH_USER_MODEL, related_name="created_feedbacksets",
                                   null=True, blank=True)

    #: The datetime when this FeedbackSet was created.
    created_datetime = models.DateTimeField(default=timezone.now)

    #: The datetime of the deadline.
    #: The first feedbackset in an AssignmentGroup
    #: (ordered by :obj:`~.FeedbackSet.created_datetime`) does not
    #: have a deadline. It inherits this from the ``first_deadline`` field
    #: of :class:`devilry.apps.core.models.assignment.Assignment`.
    deadline_datetime = models.DateTimeField(null=True, blank=True)

    #: The datetime when the feedback was published.
    #: Set when an examiner publishes the feedback for this FeedbackSet.
    #:
    #: When this is ``None``, the feedbackset is not published. This means that
    #: no comments with :obj:`.AbstractGroupComment.part_of_grading` set to ``True``
    #: is visible, the grade (extracted from points) is not visible, and this
    #: feedbackset does not count when extracting the latest/active feedback/grade
    #: for the AssignmentGroup.
    grading_published_datetime = models.DateTimeField(
        null=True,
        blank=True
    )

    #: Set when the feedbackset is published by an examiner.
    #: If this is ``None``, the feedback is not published, and
    #: the ``points`` (grade) is not available to the student.
    grading_published_by = models.ForeignKey(
        to=settings.AUTH_USER_MODEL,
        related_name="published_feedbacksets",
        null=True, blank=True
    )

    #: Points given by examiner for this feedbackset.
    #: The points on the last published FeedbackSet is the current
    #: grade for the AssignmentGroup.
    grading_points = models.PositiveIntegerField(
        null=True, blank=True
    )

    #: A :class:`django.db.models.TextField` for a gradeform filled or not filled for
    #: FeedbackSet.
    gradeform_data_json = models.TextField(
        null=False, blank=True, default=''
    )

    class Meta:
        unique_together = ('group', 'is_last_in_group')
        # index_together = (
        #     'id', 'created_datetime'
        # )

    def __unicode__(self):
        return u"{} - {} - {} - deadline: {}".format(self.group.assignment, self.feedbackset_type, self.group.long_displayname, self.deadline_datetime)

    def clean(self):
        if self.grading_published_datetime is not None and self.grading_published_by is None:
            raise ValidationError({
                'grading_published_datetime': ugettext_lazy('An assignment can not be published '
                                                            'without being published by someone.'),
            })
        if self.grading_published_datetime is not None and self.grading_points is None:
            raise ValidationError({
                'grading_published_datetime': ugettext_lazy('An assignment can not be published '
                                                            'without providing "points".'),
            })
        if self.is_last_in_group is False:
            raise ValidationError({
                'is_last_in_group': 'is_last_in_group can not be false.'
            })

    @property
    def current_deadline(self):
        if self.feedbackset_type == FeedbackSet.FEEDBACKSET_TYPE_FIRST_ATTEMPT:
            return self.group.parentnode.first_deadline or self.deadline_datetime
        return self.deadline_datetime

    def publish(self, published_by, grading_points, gradeform_data_json=''):
        """
        Publishes this FeedbackSet and comments that belongs to this feedbackset and that are
        part of the grading.

        :param published_by:
            Who published the feedbackset.

        :param grading_points:
            Points give to student(s).

        :param gradeform_data_json:
            gradeform(coming soon).

        Returns:
            True or False and an error message.

        """
        if published_by is None:
            raise ValueError
        if grading_points is None:
            raise ValueError

        current_deadline = self.current_deadline
        # if self.feedbackset_type == FeedbackSet.FEEDBACKSET_TYPE_FIRST_ATTEMPT:
        #     current_deadline = self.deadline_datetime or self.group.parentnode.first_deadline
        # elif self.feedbackset_type == FeedbackSet.FEEDBACKSET_TYPE_NEW_ATTEMPT:
        #     current_deadline = self.deadline_datetime

        if current_deadline is None:
            return False, 'Cannot publish feedback without a deadline.'

        if current_deadline > timezone.now():
            return False, 'The deadline has not expired. Feedback was saved, but not published.'

        drafted_comments = GroupComment.objects.filter(
                feedback_set=self,
                part_of_grading=True
        ).exclude_private_comments_from_other_users(
            user=published_by
        ).order_by('created_datetime')

        now = timezone.now()
        now_without_seconds = now.replace(second=0, microsecond=0)
        for modifier, draft in enumerate(drafted_comments):
            draft.publish_draft(now_without_seconds + timezone.timedelta(milliseconds=modifier))

        self.grading_points = grading_points
        self.grading_published_datetime = now_without_seconds + timezone.timedelta(
                milliseconds=drafted_comments.count() + 1)
        self.grading_published_by = published_by
        self.full_clean()
        self.save()

        return True, ''

    @property
    def gradeform_data(self):
        if self.gradeform_data_json:
            if not hasattr(self, '_gradeform_data'):
                # Store the decoded gradeform_data to avoid re-decoding the json for
                # each access. We invalidate this cache in the setter.
                self._gradeform_data = json.loads(self.gradeform_data_json)
            return self._gradeform_data
        else:
            return None

    @gradeform_data.setter
    def gradeform_data(self, gradeform_data):
        self.gradeform_data_json = json.dumps(gradeform_data)
        if hasattr(self, '_gradeform_data'):
            delattr(self, '_gradeform_data')

class GroupCommentQuerySet(AbstractGroupCommentQuerySet):
    """
    QuerySet for :class:`.GroupComment`.
    """

class GroupComment(AbstractGroupComment):
    """
    A comment made to an `AssignmentGroup`.
    """
    objects = GroupCommentQuerySet.as_manager()

    def __unicode__(self):
        return u"{} - {} - {}".format(self.feedback_set, self.user_role, self.user)

class ImageAnnotationCommentQuerySet(AbstractGroupCommentQuerySet):
    """
    QuerySet for :class:`.ImageAnnotationComment`.
    """

class ImageAnnotationComment(AbstractGroupComment):
    """
    A comment made on a file, as an annotation
    """
    objects = ImageAnnotationCommentQuerySet.as_manager()

    image = models.ForeignKey(comment_models.CommentFileImage)
    x_coordinate = models.PositiveIntegerField()
    y_coordinate = models.PositiveIntegerField()
<commit_msg>devilry_group...FeedbackSet: Added TODOs for the fields that need removal/updates.<commit_after>import json

from django.utils import timezone
from django.conf import settings
from django.core.exceptions import ValidationError
from django.db import models
from django.utils.translation import ugettext_lazy

from devilry.devilry_comment import models as comment_models
from devilry.apps.core.models import assignment_group

class AbstractGroupCommentQuerySet(models.QuerySet):
    """
    Base class for QuerySets for :class:`.AbstractGroupComment`.
    """
    def exclude_private_comments_from_other_users(self, user):
        return self.exclude(
            models.Q(visibility=AbstractGroupComment.VISIBILITY_PRIVATE) &
            ~models.Q(user=user)
        )

    def exclude_is_part_of_grading_feedbackset_unpublished(self):
        return self.exclude(
            part_of_grading=True,
            feedback_set__grading_published_datetime__isnull=True
        )

class AbstractGroupComment(comment_models.Comment):
    """
    The abstract superclass of all comments related to a delivery and feedback
    """

    #: The related feedbackset. See :class:`.FeedbackSet`.
    feedback_set = models.ForeignKey('FeedbackSet')

    #: If this is ``True``, the comment is published when the feedbackset
    #: is published. This means that this comment is part of the feedback/grading
    #: from the examiner. The same :obj:`~.AbstractGroupComment.visibility`
    #: rules apply no matter if this is ``True`` or ``False``,
    #: this only controls when the comment is published.
    part_of_grading = models.BooleanField(default=False)

    #: Comment only visible for :obj:`~devilry_comment.models.Comment.user` that created comment.
    #: Choice for :obj:`~.AbstractGroupComment.visibility`.
    VISIBILITY_PRIVATE = 'private'

    #: Comment should only be visible to examiners and admins that has
    #: access to the :obj:`~.AbstractGroupComment.feedback_set`.
    #: Choice for :obj:`~.AbstractGroupComment.visibility`.
    VISIBILITY_VISIBLE_TO_EXAMINER_AND_ADMINS = 'visible-to-examiner-and-admins'

    #: Comment should be visible to everyone that has
    #: access to the :obj:`~.AbstractGroupComment.feedback_set`.
    #: Choice for :obj:`~.AbstractGroupComment.visibility`.
    VISIBILITY_VISIBLE_TO_EVERYONE = 'visible-to-everyone'

    #: Choice list.
    #: Choices for :obj:`~.AbstractGroupComment.visibility`.
    VISIBILITY_CHOICES = [
        (VISIBILITY_PRIVATE, 'Private'),
        (VISIBILITY_VISIBLE_TO_EXAMINER_AND_ADMINS, 'Visible to examiners and admins'),
        (VISIBILITY_VISIBLE_TO_EVERYONE, 'Visible to everyone'),
    ]

    #: Sets the visibility choise of the comment.
    #: Defaults to :obj:`~.AbstractGroupComment.VISIBILITY_VISIBLE_TO_EVERYONE`.
    visibility = models.CharField(
        max_length=50,
        db_index=True,
        choices=VISIBILITY_CHOICES,
        default=VISIBILITY_VISIBLE_TO_EVERYONE,
    )

    class Meta:
        abstract = True

    def get_published_datetime(self):
        """
        Get the publishing datetime of the comment. Publishing datetime is
        the publishing time of the FeedbackSet if the comment has
        :obj:`~devilry.devilry_group.models.AbstractGroupComment.part_of_grading`
        set to True.

        Returns: Datetime.

        """
        return self.feedback_set.grading_published_datetime \
            if self.part_of_grading \
            else self.published_datetime

    def publish_draft(self, time):
        self.published_datetime = time + timezone.timedelta(milliseconds=1)
        self.full_clean()
        self.save()

class FeedbackSet(models.Model):
    """
    All comments that are given for a specific deadline (delivery and feedback) are
    linked to a feedback-set.

    If the comment has `instant_publish=True` it will be published instantly, otherwise the comments will only be
    visible once the feedbackset is published.
    All student-comments will be `instant_publish=True`, and the same applies to comments made by examiners that
    are not a part of feedback.
    """

    #: The AssignmentGroup that owns this feedbackset.
    group = models.ForeignKey(assignment_group.AssignmentGroup)

    #: Is the last feedbackset for :obj:`~.FeedbackSet.group`? Must be None or True.
    #: TODO: Remove this field - functionality is in AssignmentGroupCachedData
    is_last_in_group = models.NullBooleanField(default=True)

    #: This means the feedbackset is basically the first feedbackset.
    #: Choice for :obj:`~.FeedbackSet.feedbackset_type`.
    FEEDBACKSET_TYPE_FIRST_ATTEMPT = 'first_attempt'

    #: Is not the first feedbackset, but a new attempt.
    #: Choice for :obj:`~.FeedbackSet.feedbackset_type`
    FEEDBACKSET_TYPE_NEW_ATTEMPT = 'new_attempt'

    #: Something went wrong on grading, with this option, a new
    #: deadline should not be given to student. Student should just
    #: get notified that a new feedback was given.
    #: Choice for :obj:`~.FeedbackSet.feedbackset_type`.
    FEEDBACKSET_TYPE_RE_EDIT = 're_edit'

    #: Grading status choices for :obj:`~.FeedbackSet.feedbackset_type`.
    FEEDBACKSET_TYPE_CHOICES = [
        (FEEDBACKSET_TYPE_FIRST_ATTEMPT, 'first attempt'),
        (FEEDBACKSET_TYPE_NEW_ATTEMPT,'new attempt'),
        (FEEDBACKSET_TYPE_RE_EDIT, 're edit'),
    ]

    #: Sets the type of the feedbackset.
    #: Defaults to :obj:`~.FeedbackSet.FEEDBACKSET_TYPE_FIRST_ATTEMPT`.
    #: TODO: Remove this field? We only need the re_edit type. We can choose between keeping it, but defaulting to blank,
    #:       or replace it with an is_reedit field.
    feedbackset_type = models.CharField(
        max_length=50,
        db_index=True,
        choices=FEEDBACKSET_TYPE_CHOICES,
        default=FEEDBACKSET_TYPE_FIRST_ATTEMPT,
    )

    #: The User that created the feedbackset. Only used as metadata
    #: for superusers (for debugging).
    created_by = models.ForeignKey(settings.AUTH_USER_MODEL, related_name="created_feedbacksets",
                                   null=True, blank=True)

    #: The datetime when this FeedbackSet was created.
    created_datetime = models.DateTimeField(default=timezone.now)

    #: The datetime of the deadline.
    #: The first feedbackset in an AssignmentGroup
    #: (ordered by :obj:`~.FeedbackSet.created_datetime`) does not
    #: have a deadline. It inherits this from the ``first_deadline`` field
    #: of :class:`devilry.apps.core.models.assignment.Assignment`.
    deadline_datetime = models.DateTimeField(null=True, blank=True)

    #: The datetime when the feedback was published.
    #: Set when an examiner publishes the feedback for this FeedbackSet.
    #:
    #: When this is ``None``, the feedbackset is not published. This means that
    #: no comments with :obj:`.AbstractGroupComment.part_of_grading` set to ``True``
    #: is visible, the grade (extracted from points) is not visible, and this
    #: feedbackset does not count when extracting the latest/active feedback/grade
    #: for the AssignmentGroup.
    grading_published_datetime = models.DateTimeField(
        null=True,
        blank=True
    )

    #: Set when the feedbackset is published by an examiner.
    #: If this is ``None``, the feedback is not published, and
    #: the ``points`` (grade) is not available to the student.
    grading_published_by = models.ForeignKey(
        to=settings.AUTH_USER_MODEL,
        related_name="published_feedbacksets",
        null=True, blank=True
    )

    #: Points given by examiner for this feedbackset.
    #: The points on the last published FeedbackSet is the current
    #: grade for the AssignmentGroup.
    grading_points = models.PositiveIntegerField(
        null=True, blank=True
    )

    #: A :class:`django.db.models.TextField` for a gradeform filled or not filled for
    #: FeedbackSet.
    gradeform_data_json = models.TextField(
        null=False, blank=True, default=''
    )

    class Meta:
        unique_together = ('group', 'is_last_in_group')
        # index_together = (
        #     'id', 'created_datetime'
        # )

    def __unicode__(self):
        return u"{} - {} - {} - deadline: {}".format(self.group.assignment, self.feedbackset_type, self.group.long_displayname, self.deadline_datetime)

    def clean(self):
        if self.grading_published_datetime is not None and self.grading_published_by is None:
            raise ValidationError({
                'grading_published_datetime': ugettext_lazy('An assignment can not be published '
                                                            'without being published by someone.'),
            })
        if self.grading_published_datetime is not None and self.grading_points is None:
            raise ValidationError({
                'grading_published_datetime': ugettext_lazy('An assignment can not be published '
                                                            'without providing "points".'),
            })
        if self.is_last_in_group is False:
            raise ValidationError({
                'is_last_in_group': 'is_last_in_group can not be false.'
            })

    @property
    def current_deadline(self):
        if self.feedbackset_type == FeedbackSet.FEEDBACKSET_TYPE_FIRST_ATTEMPT:
            return self.group.parentnode.first_deadline or self.deadline_datetime
        return self.deadline_datetime

    def publish(self, published_by, grading_points, gradeform_data_json=''):
        """
        Publishes this FeedbackSet and comments that belongs to this feedbackset and that are
        part of the grading.

        :param published_by:
            Who published the feedbackset.

        :param grading_points:
            Points give to student(s).

        :param gradeform_data_json:
            gradeform(coming soon).

        Returns:
            True or False and an error message.

        """
        if published_by is None:
            raise ValueError
        if grading_points is None:
            raise ValueError

        current_deadline = self.current_deadline
        # if self.feedbackset_type == FeedbackSet.FEEDBACKSET_TYPE_FIRST_ATTEMPT:
        #     current_deadline = self.deadline_datetime or self.group.parentnode.first_deadline
        # elif self.feedbackset_type == FeedbackSet.FEEDBACKSET_TYPE_NEW_ATTEMPT:
        #     current_deadline = self.deadline_datetime

        if current_deadline is None:
            return False, 'Cannot publish feedback without a deadline.'

        if current_deadline > timezone.now():
            return False, 'The deadline has not expired. Feedback was saved, but not published.'

        drafted_comments = GroupComment.objects.filter(
                feedback_set=self,
                part_of_grading=True
        ).exclude_private_comments_from_other_users(
            user=published_by
        ).order_by('created_datetime')

        now = timezone.now()
        now_without_seconds = now.replace(second=0, microsecond=0)
        for modifier, draft in enumerate(drafted_comments):
            draft.publish_draft(now_without_seconds + timezone.timedelta(milliseconds=modifier))

        self.grading_points = grading_points
        self.grading_published_datetime = now_without_seconds + timezone.timedelta(
                milliseconds=drafted_comments.count() + 1)
        self.grading_published_by = published_by
        self.full_clean()
        self.save()

        return True, ''

    @property
    def gradeform_data(self):
        if self.gradeform_data_json:
            if not hasattr(self, '_gradeform_data'):
                # Store the decoded gradeform_data to avoid re-decoding the json for
                # each access. We invalidate this cache in the setter.
                self._gradeform_data = json.loads(self.gradeform_data_json)
            return self._gradeform_data
        else:
            return None

    @gradeform_data.setter
    def gradeform_data(self, gradeform_data):
        self.gradeform_data_json = json.dumps(gradeform_data)
        if hasattr(self, '_gradeform_data'):
            delattr(self, '_gradeform_data')

class GroupCommentQuerySet(AbstractGroupCommentQuerySet):
    """
    QuerySet for :class:`.GroupComment`.
    """

class GroupComment(AbstractGroupComment):
    """
    A comment made to an `AssignmentGroup`.
    """
    objects = GroupCommentQuerySet.as_manager()

    def __unicode__(self):
        return u"{} - {} - {}".format(self.feedback_set, self.user_role, self.user)

class ImageAnnotationCommentQuerySet(AbstractGroupCommentQuerySet):
    """
    QuerySet for :class:`.ImageAnnotationComment`.
    """

class ImageAnnotationComment(AbstractGroupComment):
    """
    A comment made on a file, as an annotation
    """
    objects = ImageAnnotationCommentQuerySet.as_manager()

    image = models.ForeignKey(comment_models.CommentFileImage)
    x_coordinate = models.PositiveIntegerField()
    y_coordinate = models.PositiveIntegerField()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2016 Yelp Inc.
#
# 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 mock
import pytest

from .helper import broker_range
from kafka_utils.kafka_cluster_manager.cluster_info \
    .error import BrokerDecommissionError
from kafka_utils.kafka_cluster_manager.cluster_info \
    .partition_count_balancer import PartitionCountBalancer

class TestClusterBalancer(object):

    @pytest.fixture(params=[PartitionCountBalancer])
    def create_balancer(self, request):
        def build_balancer(cluster_topology, **kwargs):
            args = mock.Mock(**kwargs)
            return request.param(cluster_topology, args)
        return build_balancer

    def assert_rg_balanced_partition(self, ct, partition, opt_cnt, extra_cnt=0):
        """Verify that the partition replicas are balanced across
        replication-groups.
        """
        for rg in ct.rgs.itervalues():
            replica_cnt_rg = rg.count_replica(partition)

            # Verify evenly-balanced partition
            assert replica_cnt_rg == opt_cnt or\
                replica_cnt_rg == opt_cnt + extra_cnt

    def test_broker_decommission(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = {
            (u'T0', 0): ['0', '2'],
            (u'T0', 1): ['0', '3'],
            (u'T1', 0): ['0', '5'],
        }
        ct = create_cluster_topology(assignment)
        partitions_count = len(ct.partitions)

        # should move all partitions from broker 0 to either 1 or 4 because they
        # are in the same replication group and empty.
        cb = create_balancer(ct)
        cb.decommission_brokers(['0'])

        # Here we just care that broker 0 is empty and partitions count didn't
        # change
        assert len(ct.partitions) == partitions_count
        assert ct.brokers['0'].empty()

    def test_broker_decommission_many(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = {
            (u'T0', 0): ['0', '2'],
            (u'T0', 1): ['0', '3'],
            (u'T1', 0): ['0', '5'],
            (u'T1', 1): ['1', '5'],
            (u'T1', 2): ['1', '5'],
            (u'T2', 0): ['0', '3'],
        }
        ct = create_cluster_topology(assignment)
        partitions_count = len(ct.partitions)

        cb = create_balancer(ct)
        cb.decommission_brokers(['0', '1', '3'])

        assert len(ct.partitions) == partitions_count
        assert ct.brokers['0'].empty()
        assert ct.brokers['1'].empty()
        assert ct.brokers['3'].empty()

    def test_broker_decommission_force(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = {
            (u'T0', 0): ['0', '1', '2'],
            (u'T0', 1): ['0', '1', '2'],
            (u'T1', 0): ['0', '1'],
            (u'T1', 1): ['1', '4'],
            (u'T1', 2): ['1', '4'],
            (u'T2', 0): ['0', '3'],
        }
        # r1 b0 t00, t01, t10, t20
        # r1 b1 t00, t01, t10, t11, t12
        # r2 b2 t00, t01
        # r2 b3 t20
        # r3 b4 t11, t12
        brokers_rg = {'0': 'rg1', '1': 'rg1', '2': 'rg2', '3': 'rg2', '4': 'r3'}
        ct = create_cluster_topology(
            assignment,
            brokers_rg,
            lambda x: x.metadata,  # The value of the broker dict is the metadata attribute
        )
        partitions_count = len(ct.partitions)

        cb = create_balancer(ct)
        cb.decommission_brokers(['0'])

        assert len(ct.partitions) == partitions_count
        assert ct.brokers['0'].empty()

    def test_broker_decommission_empty_replication_group(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = {
            (u'T0', 0): ['0', '1', '2'],
            (u'T0', 1): ['0', '1', '2'],
            (u'T1', 0): ['0', '1'],
            (u'T1', 1): ['1', '2'],
            (u'T1', 2): ['1', '2'],
            (u'T2', 0): ['0', '2'],
        }
        # r0 b0 t00, t01, t10, t20
        # r1 b1 t00, t01, t10, t11, t12
        # __ b2 t00, t01, t11, t12, t20 let's assume b2 is down and there is no
        # metadata for it (it was in r2 before the failure)
        # r2 b3 empty this broker came up to replace b2
        brokers_rg = {'0': 'rg0', '1': 'rg1', '3': 'rg2'}  # NOTE: b2 is not in this list

        ct = create_cluster_topology(
            assignment,
            brokers_rg,
            lambda x: x.metadata,  # The value of the broker dict is the metadata attribute
        )
        partitions_count = len(ct.partitions)

        cb = create_balancer(ct)
        cb.decommission_brokers(['2'])

        assert len(ct.partitions) == partitions_count
        assert ct.brokers['2'].empty()

    def test_elect_source_replication_group(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        # Sample assignment with 3 replication groups
        # with replica-count as as per map :-
        # broker_rg = {0: 'rg1', 1: 'rg1', 2: 'rg2', 3: 'rg2', 4: 'rg1', 5: 'rg3', 6: 'rg4'}
        # rg-id:    (brokers), count
        # rg1:      (0, 2, 4) = 3
        # rg2:      (1, 3) = 2
        # rg3:      (5) = 1
        # rg4:      (6) = 1
        # rg1 and rg2 are over-replicated and rg3 being under-replicated
        # source-replication-group should be rg1 having the highest replicas
        p1_info = ((u'T0', 0), ['0', '1', '2', '3', '4', '5', '6'])
        assignment = dict([p1_info])
        ct = create_cluster_topology(assignment, broker_range(7))
        p1 = ct.partitions[p1_info[0]]
        cb = create_balancer(ct)
        # Case-1: rg's have only 1 unique max replica count
        # 'rg1' and 'rg2' are over-replicated groups
        over_replicated_rgs = [ct.rgs['rg1'], ct.rgs['rg2']]

        # Get source-replication group
        rg_source = cb._elect_source_replication_group(over_replicated_rgs, p1)

        # Since, 'rg1' has more replicas i.e. 3, it should be selected
        assert rg_source.id == 'rg1'

    def test_elect_dest_replication_group(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        # Sample assignment with 3 replication groups
        # with replica-count as as per map
        # broker_rg: {0: 'rg1', 1: 'rg1', 2: 'rg2', 3: 'rg2', 4: 'rg1', 5: 'rg3'}
        # rg-id: (brokers), count
        # rg1: (0, 2, 4) = 3
        # rg2: (1, 3) = 2
        # rg3: (5) = 1
        # rg1 and rg2 are over-replicated and rg3 being under-replicated
        # source-replication-group should be rg1 having the highest replicas
        p1_info = ((u'T0', 0), ['0', '1', '2', '3', '4', '5', '6'])
        assignment = dict([p1_info])
        ct = create_cluster_topology(assignment, broker_range(7))
        p1 = ct.partitions[p1_info[0]]
        cb = create_balancer(ct)
        # Case 1: rg_source = 'rg1', find destination-replica
        rg_source = ct.rgs['rg1']
        under_replicated_rgs = [ct.rgs['rg3'], ct.rgs['rg4']]
        # Get destination-replication-group for partition: p1
        rg_dest = cb._elect_dest_replication_group(
            rg_source.count_replica(p1),
            under_replicated_rgs,
            p1,
        )

        # Dest-replica can be either 'rg3' or 'rg4' with replica-count 1
        assert rg_dest.id in ['rg3', 'rg4']

        # Case 2: rg-source == 'rg2': No destination group found
        rg_source = ct.rgs['rg2']
        # Get destination-replication-group for partition: p1
        rg_dest = cb._elect_dest_replication_group(
            rg_source.count_replica(p1),
            under_replicated_rgs,
            p1,
        )

        # Since none of under-replicated-groups (rg3, and rg4) have lower
        # 2-1=0 replicas for the given partition p1
        # No eligible dest-group is there where partition can be sent to
        assert rg_dest is None

    def test_rebalance_partition_replicas_imbalanced_case1(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        # Test imbalanced partitions for below cases
        # Note: In initial-assignment, all partitions with id-1 are 'imbalanced'
        ct = create_cluster_topology()
        cb = create_balancer(ct)
        # CASE 1: repl-factor % rg-count == 0
        # (1a): repl-factor == rg-count
        # p1: replicas: ('T0', 1): [2,3]
        p1 = ct.partitions[('T0', 1)]
        # rg-imbalanced p1
        opt_cnt = 1    # 2/2

        cb._rebalance_partition_replicas(p1)

        # Verify partition is rg-balanced
        self.assert_rg_balanced_partition(ct, p1, opt_cnt)

        # (1b):  repl-count % rg-count == 0 and repl-count > rg-count
        # p1: replicas: ('T1',1): [0,1,2,4]
        p1 = ct.partitions[('T1', 1)]
        # Assert originally-imbalanced p1
        opt_cnt = 2    # 4/2
        cb._rebalance_partition_replicas(p1)

        # Verify partition is rg-balanced
        self.assert_rg_balanced_partition(ct, p1, opt_cnt)

    def test_rebalance_partition_replicas_imbalanced_case2(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        ct = create_cluster_topology()
        cb = create_balancer(ct)
        # CASE 2: repl-factor % rg-count > 0
        # p1: replicas ('T3', 1): [0,1,4]
        p1 = ct.partitions[('T3', 1)]
        # rg-imbalanced p1
        opt_cnt = 1    # 3/2
        extra_cnt = 1  # 3%2
        cb._rebalance_partition_replicas(p1)

        # Verify partition is now rg-balanced
        self.assert_rg_balanced_partition(ct, p1, opt_cnt, extra_cnt)

    def test_rebalance_partition_replicas_balanced(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        # Test already balanced partitions in given example for different cases
        # Analyze Cases 1a, 1b
        ct = create_cluster_topology()
        cb = create_balancer(ct)
        # CASE 1: repl-factor % rg-count == 0
        # (1a): repl-factor == rg-count
        # p1: replicas: ('T0', 0): [1,2]
        p1 = ct.partitions[('T0', 0)]
        opt_cnt = 1    # 2/2
        self.assert_rg_balanced_partition(ct, p1, opt_cnt)
        cb._rebalance_partition_replicas(p1)

        # Verify no change in replicas after rebalancing
        p1_old_replicas = p1.replicas
        cb._rebalance_partition(p1)
        assert p1_old_replicas == p1.replicas

        # (1b):  repl-count % rg-count == 0 and repl-count > rg-count
        # p2: replicas: ('T1',0): [0,1,2,3]
        p2 = ct.partitions[('T1', 0)]
        opt_cnt = 2    # 4/2
        self.assert_rg_balanced_partition(ct, p2, opt_cnt)

        # Verify no change in replicas after rebalancing
        p2_old_replicas = p2.replicas
        cb._rebalance_partition(p2)
        assert p2_old_replicas == p2.replicas

    def test_broker_decommission_error(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = {
            (u'T1', 0): ['0', '1', '2', '3', '4'],
            (u'T1', 1): ['0', '1', '2', '4'],
            (u'T2', 1): ['0', '1', '2', '4'],
        }
        ct = create_cluster_topology(assignment)

        with pytest.raises(BrokerDecommissionError):
            cb = create_balancer(ct)
            cb.decommission_brokers(['0'])

    def test_add_replica(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = dict([((u'T1', 0), ['1', '3'])])
        ct = create_cluster_topology(assignment, broker_range(6))
        partition = ct.partitions[(u'T1', 0)]

        cb = create_balancer(ct)
        cb.add_replica(partition.name, count=3)

        assert partition.replication_factor == 5
        assert sum(rg.count_replica(partition) for rg in ct.rgs.values()) == 5

    def test_remove_replica(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = dict([((u'T1', 0), ['0', '1', '2', '3', '5'])])
        ct = create_cluster_topology(assignment, broker_range(6))
        partition = ct.partitions[(u'T1', 0)]
        osr_broker_ids = ['1', '2']

        cb = create_balancer(ct)
        cb.remove_replica(partition.name, osr_broker_ids, count=3)

        assert partition.replication_factor == 2
        assert sum(rg.count_replica(partition) for rg in ct.rgs.values()) == 2

    def test_remove_replica_prioritize_osr(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = dict([((u'T1', 0), ['0', '1', '2', '3', '5'])])
        ct = create_cluster_topology(assignment, broker_range(6))
        partition = ct.partitions[(u'T1', 0)]
        osr_broker_ids = ['1', '2']

        cb = create_balancer(ct)
        cb.remove_replica(partition.name, osr_broker_ids, count=2)

        assert set(b.id for b in partition.replicas) == set(['0', '3', '5'])
<commit_msg>Enable testing for GeneticBalancer<commit_after># -*- coding: utf-8 -*-
# Copyright 2016 Yelp Inc.
#
# 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 mock
import pytest

from .helper import broker_range
from kafka_utils.kafka_cluster_manager.cluster_info \
    .error import BrokerDecommissionError
from kafka_utils.kafka_cluster_manager.cluster_info \
    .genetic_balancer import GeneticBalancer
from kafka_utils.kafka_cluster_manager.cluster_info \
    .partition_count_balancer import PartitionCountBalancer

class TestClusterBalancer(object):

    @pytest.fixture(params=[PartitionCountBalancer, GeneticBalancer])
    def create_balancer(self, request):
        def build_balancer(cluster_topology, **kwargs):
            args = mock.Mock(**kwargs)
            return request.param(cluster_topology, args)
        return build_balancer

    def assert_rg_balanced_partition(self, ct, partition, opt_cnt, extra_cnt=0):
        """Verify that the partition replicas are balanced across
        replication-groups.
        """
        for rg in ct.rgs.itervalues():
            replica_cnt_rg = rg.count_replica(partition)

            # Verify evenly-balanced partition
            assert replica_cnt_rg == opt_cnt or\
                replica_cnt_rg == opt_cnt + extra_cnt

    def test_broker_decommission(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = {
            (u'T0', 0): ['0', '2'],
            (u'T0', 1): ['0', '3'],
            (u'T1', 0): ['0', '5'],
        }
        ct = create_cluster_topology(assignment)
        partitions_count = len(ct.partitions)

        # should move all partitions from broker 0 to either 1 or 4 because they
        # are in the same replication group and empty.
        cb = create_balancer(ct)
        cb.decommission_brokers(['0'])

        # Here we just care that broker 0 is empty and partitions count didn't
        # change
        assert len(ct.partitions) == partitions_count
        assert ct.brokers['0'].empty()

    def test_broker_decommission_many(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = {
            (u'T0', 0): ['0', '2'],
            (u'T0', 1): ['0', '3'],
            (u'T1', 0): ['0', '5'],
            (u'T1', 1): ['1', '5'],
            (u'T1', 2): ['1', '5'],
            (u'T2', 0): ['0', '3'],
        }
        ct = create_cluster_topology(assignment)
        partitions_count = len(ct.partitions)

        cb = create_balancer(ct)
        cb.decommission_brokers(['0', '1', '3'])

        assert len(ct.partitions) == partitions_count
        assert ct.brokers['0'].empty()
        assert ct.brokers['1'].empty()
        assert ct.brokers['3'].empty()

    def test_broker_decommission_force(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = {
            (u'T0', 0): ['0', '1', '2'],
            (u'T0', 1): ['0', '1', '2'],
            (u'T1', 0): ['0', '1'],
            (u'T1', 1): ['1', '4'],
            (u'T1', 2): ['1', '4'],
            (u'T2', 0): ['0', '3'],
        }
        # r1 b0 t00, t01, t10, t20
        # r1 b1 t00, t01, t10, t11, t12
        # r2 b2 t00, t01
        # r2 b3 t20
        # r3 b4 t11, t12
        brokers_rg = {'0': 'rg1', '1': 'rg1', '2': 'rg2', '3': 'rg2', '4': 'r3'}
        ct = create_cluster_topology(
            assignment,
            brokers_rg,
            lambda x: x.metadata,  # The value of the broker dict is the metadata attribute
        )
        partitions_count = len(ct.partitions)

        cb = create_balancer(ct)
        cb.decommission_brokers(['0'])

        assert len(ct.partitions) == partitions_count
        assert ct.brokers['0'].empty()

    def test_broker_decommission_empty_replication_group(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = {
            (u'T0', 0): ['0', '1', '2'],
            (u'T0', 1): ['0', '1', '2'],
            (u'T1', 0): ['0', '1'],
            (u'T1', 1): ['1', '2'],
            (u'T1', 2): ['1', '2'],
            (u'T2', 0): ['0', '2'],
        }
        # r0 b0 t00, t01, t10, t20
        # r1 b1 t00, t01, t10, t11, t12
        # __ b2 t00, t01, t11, t12, t20 let's assume b2 is down and there is no
        # metadata for it (it was in r2 before the failure)
        # r2 b3 empty this broker came up to replace b2
        brokers_rg = {'0': 'rg0', '1': 'rg1', '3': 'rg2'}  # NOTE: b2 is not in this list

        ct = create_cluster_topology(
            assignment,
            brokers_rg,
            lambda x: x.metadata,  # The value of the broker dict is the metadata attribute
        )
        partitions_count = len(ct.partitions)

        cb = create_balancer(ct)
        cb.decommission_brokers(['2'])

        assert len(ct.partitions) == partitions_count
        assert ct.brokers['2'].empty()

    def test_elect_source_replication_group(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        # Sample assignment with 3 replication groups
        # with replica-count as as per map :-
        # broker_rg = {0: 'rg1', 1: 'rg1', 2: 'rg2', 3: 'rg2', 4: 'rg1', 5: 'rg3', 6: 'rg4'}
        # rg-id:    (brokers), count
        # rg1:      (0, 2, 4) = 3
        # rg2:      (1, 3) = 2
        # rg3:      (5) = 1
        # rg4:      (6) = 1
        # rg1 and rg2 are over-replicated and rg3 being under-replicated
        # source-replication-group should be rg1 having the highest replicas
        p1_info = ((u'T0', 0), ['0', '1', '2', '3', '4', '5', '6'])
        assignment = dict([p1_info])
        ct = create_cluster_topology(assignment, broker_range(7))
        p1 = ct.partitions[p1_info[0]]
        cb = create_balancer(ct)
        # Case-1: rg's have only 1 unique max replica count
        # 'rg1' and 'rg2' are over-replicated groups
        over_replicated_rgs = [ct.rgs['rg1'], ct.rgs['rg2']]

        # Get source-replication group
        rg_source = cb._elect_source_replication_group(over_replicated_rgs, p1)

        # Since, 'rg1' has more replicas i.e. 3, it should be selected
        assert rg_source.id == 'rg1'

    def test_elect_dest_replication_group(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        # Sample assignment with 3 replication groups
        # with replica-count as as per map
        # broker_rg: {0: 'rg1', 1: 'rg1', 2: 'rg2', 3: 'rg2', 4: 'rg1', 5: 'rg3'}
        # rg-id: (brokers), count
        # rg1: (0, 2, 4) = 3
        # rg2: (1, 3) = 2
        # rg3: (5) = 1
        # rg1 and rg2 are over-replicated and rg3 being under-replicated
        # source-replication-group should be rg1 having the highest replicas
        p1_info = ((u'T0', 0), ['0', '1', '2', '3', '4', '5', '6'])
        assignment = dict([p1_info])
        ct = create_cluster_topology(assignment, broker_range(7))
        p1 = ct.partitions[p1_info[0]]
        cb = create_balancer(ct)
        # Case 1: rg_source = 'rg1', find destination-replica
        rg_source = ct.rgs['rg1']
        under_replicated_rgs = [ct.rgs['rg3'], ct.rgs['rg4']]
        # Get destination-replication-group for partition: p1
        rg_dest = cb._elect_dest_replication_group(
            rg_source.count_replica(p1),
            under_replicated_rgs,
            p1,
        )

        # Dest-replica can be either 'rg3' or 'rg4' with replica-count 1
        assert rg_dest.id in ['rg3', 'rg4']

        # Case 2: rg-source == 'rg2': No destination group found
        rg_source = ct.rgs['rg2']
        # Get destination-replication-group for partition: p1
        rg_dest = cb._elect_dest_replication_group(
            rg_source.count_replica(p1),
            under_replicated_rgs,
            p1,
        )

        # Since none of under-replicated-groups (rg3, and rg4) have lower
        # 2-1=0 replicas for the given partition p1
        # No eligible dest-group is there where partition can be sent to
        assert rg_dest is None

    def test_rebalance_partition_replicas_imbalanced_case1(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        # Test imbalanced partitions for below cases
        # Note: In initial-assignment, all partitions with id-1 are 'imbalanced'
        ct = create_cluster_topology()
        cb = create_balancer(ct)
        # CASE 1: repl-factor % rg-count == 0
        # (1a): repl-factor == rg-count
        # p1: replicas: ('T0', 1): [2,3]
        p1 = ct.partitions[('T0', 1)]
        # rg-imbalanced p1
        opt_cnt = 1    # 2/2

        cb._rebalance_partition_replicas(p1)

        # Verify partition is rg-balanced
        self.assert_rg_balanced_partition(ct, p1, opt_cnt)

        # (1b):  repl-count % rg-count == 0 and repl-count > rg-count
        # p1: replicas: ('T1',1): [0,1,2,4]
        p1 = ct.partitions[('T1', 1)]
        # Assert originally-imbalanced p1
        opt_cnt = 2    # 4/2
        cb._rebalance_partition_replicas(p1)

        # Verify partition is rg-balanced
        self.assert_rg_balanced_partition(ct, p1, opt_cnt)

    def test_rebalance_partition_replicas_imbalanced_case2(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        ct = create_cluster_topology()
        cb = create_balancer(ct)
        # CASE 2: repl-factor % rg-count > 0
        # p1: replicas ('T3', 1): [0,1,4]
        p1 = ct.partitions[('T3', 1)]
        # rg-imbalanced p1
        opt_cnt = 1    # 3/2
        extra_cnt = 1  # 3%2
        cb._rebalance_partition_replicas(p1)

        # Verify partition is now rg-balanced
        self.assert_rg_balanced_partition(ct, p1, opt_cnt, extra_cnt)

    def test_rebalance_partition_replicas_balanced(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        # Test already balanced partitions in given example for different cases
        # Analyze Cases 1a, 1b
        ct = create_cluster_topology()
        cb = create_balancer(ct)
        # CASE 1: repl-factor % rg-count == 0
        # (1a): repl-factor == rg-count
        # p1: replicas: ('T0', 0): [1,2]
        p1 = ct.partitions[('T0', 0)]
        opt_cnt = 1    # 2/2
        self.assert_rg_balanced_partition(ct, p1, opt_cnt)
        cb._rebalance_partition_replicas(p1)

        # Verify no change in replicas after rebalancing
        p1_old_replicas = p1.replicas
        cb._rebalance_partition(p1)
        assert p1_old_replicas == p1.replicas

        # (1b):  repl-count % rg-count == 0 and repl-count > rg-count
        # p2: replicas: ('T1',0): [0,1,2,3]
        p2 = ct.partitions[('T1', 0)]
        opt_cnt = 2    # 4/2
        self.assert_rg_balanced_partition(ct, p2, opt_cnt)

        # Verify no change in replicas after rebalancing
        p2_old_replicas = p2.replicas
        cb._rebalance_partition(p2)
        assert p2_old_replicas == p2.replicas

    def test_broker_decommission_error(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = {
            (u'T1', 0): ['0', '1', '2', '3', '4'],
            (u'T1', 1): ['0', '1', '2', '4'],
            (u'T2', 1): ['0', '1', '2', '4'],
        }
        ct = create_cluster_topology(assignment)

        with pytest.raises(BrokerDecommissionError):
            cb = create_balancer(ct)
            cb.decommission_brokers(['0'])

    def test_add_replica(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = dict([((u'T1', 0), ['1', '3'])])
        ct = create_cluster_topology(assignment, broker_range(6))
        partition = ct.partitions[(u'T1', 0)]

        cb = create_balancer(ct)
        cb.add_replica(partition.name, count=3)

        assert partition.replication_factor == 5
        assert sum(rg.count_replica(partition) for rg in ct.rgs.values()) == 5

    def test_remove_replica(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = dict([((u'T1', 0), ['0', '1', '2', '3', '5'])])
        ct = create_cluster_topology(assignment, broker_range(6))
        partition = ct.partitions[(u'T1', 0)]
        osr_broker_ids = ['1', '2']

        cb = create_balancer(ct)
        cb.remove_replica(partition.name, osr_broker_ids, count=3)

        assert partition.replication_factor == 2
        assert sum(rg.count_replica(partition) for rg in ct.rgs.values()) == 2

    def test_remove_replica_prioritize_osr(
            self,
            create_balancer,
            create_cluster_topology,
    ):
        assignment = dict([((u'T1', 0), ['0', '1', '2', '3', '5'])])
        ct = create_cluster_topology(assignment, broker_range(6))
        partition = ct.partitions[(u'T1', 0)]
        osr_broker_ids = ['1', '2']

        cb = create_balancer(ct)
        cb.remove_replica(partition.name, osr_broker_ids, count=2)

        assert set(b.id for b in partition.replicas) == set(['0', '3', '5'])
<|endoftext|>


<commit_before>import random

def select_movie(movie_list):
    """Selects a movie from the database"""
    return random.choice(movie_list)
<commit_msg>Better randomness through using entropy from the system<commit_after>import random

def select_movie(movie_list):
    """Selects a movie from the database"""
    return random.SystemRandom().choice(movie_list)
<|endoftext|>


<commit_before>__author__ = 'nlw'

from dipper.utils.CurieUtil import CurieUtil
from dipper.utils.GraphUtils import GraphUtils
from dipper.models.Assoc import Assoc
from dipper import curie_map

class OrthologyAssoc(Assoc):

    ortho_rel = {
        'orthologous': 'RO:HOM0000017',  # in orthology relationship with
        'least_diverged_orthologous': 'RO:HOM0000020',  # in 1 to 1 orthology relationship with
        'homologous': 'RO:HOM0000019',  # in 1 to 1 homology relationship with
        'paralogous': 'RO:HOM0000011',  # in paralogy relationship with (generic)
        'in_paralogous': 'RO:HOM0000023',  # in in-paralogy relationship with
        'ohnologous': 'RO:HOM0000022',  # in ohnology relationship with
        'xenologous': 'RO:HOM0000018',  # in xenology relationship with
        'has_member': 'RO:0002351'
        }

    def __init__(self, assoc_id, gene1, gene2, pub, evidence_code):
        super().__init__()
        self.cu = CurieUtil(curie_map.get())
        self.gu = GraphUtils(curie_map.get())
        self.object_properties.update(self.ortho_rel)
        self.properties.update(self.ortho_rel)
        self.annot_id = assoc_id
        self.gene1 = gene1
        self.gene2 = gene2
        self.pub_id = pub
        self.evidence = evidence_code
        self.rel = self.properties['orthologous']  # default
        self.pub_list = None

        self.setSubject(gene1)
        self.setObject(gene2)

        return

    def addOrthologyAssociationToGraph(self, g):

        self.addAssociationToGraph(g)

        return

    def addGeneFamilyToGraph(self, g, family_id):
        """
        Make an association between a group of genes and some grouping class.
        We make the assumption that the genes in the association are part of the supplied
        family_id, and that the genes have already been declared as classes elsewhere.
        The family_id is added as an individual of type DATA:gene_family.

        Triples:
        <family_id> a DATA:gene_family
        <family_id> RO:has_member <gene1>
        <family_id> RO:has_member <gene2>

        :param family_id:
        :param g: the graph to modify
        :return:
        """
        gene_family = 'DATA:3148'  # http://edamontology.org/data_3148

        # make the assumption that the genes have already been added as classes previously
        self.gu.addIndividualToGraph(g, family_id, None, gene_family)

        # add each gene to the family
        self.gu.addMember(g, family_id, self.gene1)
        self.gu.addMember(g, family_id, self.gene2)

        return
<commit_msg>add Othology terms of gene_family<commit_after>__author__ = 'nlw'

from dipper.utils.CurieUtil import CurieUtil
from dipper.utils.GraphUtils import GraphUtils
from dipper.models.Assoc import Assoc
from dipper import curie_map

class OrthologyAssoc(Assoc):

    ortho_rel = {
        'orthologous': 'RO:HOM0000017',  # in orthology relationship with
        'least_diverged_orthologous': 'RO:HOM0000020',  # in 1 to 1 orthology relationship with
        'homologous': 'RO:HOM0000019',  # in 1 to 1 homology relationship with
        'paralogous': 'RO:HOM0000011',  # in paralogy relationship with (generic)
        'in_paralogous': 'RO:HOM0000023',  # in in-paralogy relationship with
        'ohnologous': 'RO:HOM0000022',  # in ohnology relationship with
        'xenologous': 'RO:HOM0000018',  # in xenology relationship with
        'has_member': 'RO:0002351'
        }

    terms = {
        'gene_family': 'DATA:3148'  # http://edamontology.org/data_3148
    }

    def __init__(self, assoc_id, gene1, gene2, pub, evidence_code):
        super().__init__()
        self.cu = CurieUtil(curie_map.get())
        self.gu = GraphUtils(curie_map.get())
        self.object_properties.update(self.ortho_rel)
        self.properties.update(self.ortho_rel)
        self.annot_id = assoc_id
        self.gene1 = gene1
        self.gene2 = gene2
        self.pub_id = pub
        self.evidence = evidence_code
        self.rel = self.properties['orthologous']  # default
        self.pub_list = None

        self.setSubject(gene1)
        self.setObject(gene2)

        return

    def addOrthologyAssociationToGraph(self, g):

        self.addAssociationToGraph(g)

        return

    def addGeneFamilyToGraph(self, g, family_id):
        """
        Make an association between a group of genes and some grouping class.
        We make the assumption that the genes in the association are part of the supplied
        family_id, and that the genes have already been declared as classes elsewhere.
        The family_id is added as an individual of type DATA:gene_family.

        Triples:
        <family_id> a DATA:gene_family
        <family_id> RO:has_member <gene1>
        <family_id> RO:has_member <gene2>

        :param family_id:
        :param g: the graph to modify
        :return:
        """
        gene_family = self.terms['gene_family']  # http://edamontology.org/data_3148

        # make the assumption that the genes have already been added as classes previously
        self.gu.addIndividualToGraph(g, family_id, None, gene_family)

        # add each gene to the family
        self.gu.addMember(g, family_id, self.gene1)
        self.gu.addMember(g, family_id, self.gene2)

        return
<|endoftext|>