<commit_before>from decimal import Decimal
from django.core.serializers import deserialize, serialize
from django.core.serializers.base import DeserializationError
from django.db import models
from django.test import TestCase
from django.utils import simplejson as json

from .fields import JSONField, JSONCharField
from django.forms.util import ValidationError

from collections import OrderedDict

class JsonModel(models.Model):
    json = JSONField()
    default_json = JSONField(default={"check":12})
    complex_default_json = JSONField(default=[{"checkcheck": 1212}])

class JsonCharModel(models.Model):
    json = JSONCharField(max_length=100)
    default_json = JSONCharField(max_length=100, default={"check":34})

class ComplexEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, complex):
            return {
                '__complex__': True,
                'real': obj.real,
                'imag': obj.imag,
            }

        return json.JSONEncoder.default(self, obj)

def as_complex(dct):
    if '__complex__' in dct:
        return complex(dct['real'], dct['imag'])
    return dct

class JSONModelCustomEncoders(models.Model):
    # A JSON field that can store complex numbers
    json = JSONField(
        dump_kwargs={'cls': ComplexEncoder, "indent": 4},
        load_kwargs={'object_hook': as_complex},
    )

class JSONFieldTest(TestCase):
    """JSONField Wrapper Tests"""

    json_model = JsonModel

    def test_json_field_create(self):
        """Test saving a JSON object in our JSONField"""
        json_obj = {
            "item_1": "this is a json blah",
            "blergh": "hey, hey, hey"}

        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_string_in_json_field(self):
        """Test saving an ordinary Python string in our JSONField"""
        json_obj = 'blah blah'
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_float_in_json_field(self):
        """Test saving a Python float in our JSONField"""
        json_obj = 1.23
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_int_in_json_field(self):
        """Test saving a Python integer in our JSONField"""
        json_obj = 1234567
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_decimal_in_json_field(self):
        """Test saving a Python Decimal in our JSONField"""
        json_obj = Decimal(12.34)
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        # here we must know to convert the returned string back to Decimal,
        # since json does not support that format
        self.assertEqual(Decimal(new_obj.json), json_obj)

    def test_json_field_modify(self):
        """Test modifying a JSON object in our JSONField"""
        json_obj_1 = {'a': 1, 'b': 2}
        json_obj_2 = {'a': 3, 'b': 4}

        obj = self.json_model.objects.create(json=json_obj_1)
        self.assertEqual(obj.json, json_obj_1)
        obj.json = json_obj_2

        self.assertEqual(obj.json, json_obj_2)
        obj.save()
        self.assertEqual(obj.json, json_obj_2)

        self.assertTrue(obj)

    def test_json_field_load(self):
        """Test loading a JSON object from the DB"""
        json_obj_1 = {'a': 1, 'b': 2}
        obj = self.json_model.objects.create(json=json_obj_1)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj_1)

    def test_json_list(self):
        """Test storing a JSON list"""
        json_obj = ["my", "list", "of", 1, "objs", {"hello": "there"}]

        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)
        self.assertEqual(new_obj.json, json_obj)

    def test_empty_objects(self):
        """Test storing empty objects"""
        for json_obj in [{}, [], 0, '', False]:
            obj = self.json_model.objects.create(json=json_obj)
            new_obj = self.json_model.objects.get(id=obj.id)
            self.assertEqual(json_obj, obj.json)
            self.assertEqual(json_obj, new_obj.json)

    def test_custom_encoder(self):
        """Test encoder_cls and object_hook"""
        value = 1 + 3j  # A complex number

        obj = JSONModelCustomEncoders.objects.create(json=value)
        new_obj = JSONModelCustomEncoders.objects.get(pk=obj.pk)
        self.assertEqual(value, new_obj.json)

    def test_django_serializers(self):
        """Test serializing/deserializing jsonfield data"""
        for json_obj in [{}, [], 0, '', False, {'key': 'value', 'num': 42,
                                                'ary': list(range(5)),
                                                'dict': {'k': 'v'}}]:
            obj = self.json_model.objects.create(json=json_obj)
            new_obj = self.json_model.objects.get(id=obj.id)

        queryset = self.json_model.objects.all()
        ser = serialize('json', queryset)
        for dobj in deserialize('json', ser):
            obj = dobj.object
            pulled = self.json_model.objects.get(id=obj.pk)
            self.assertEqual(obj.json, pulled.json)

    def test_default_parameters(self):
        """Test providing a default value to the model"""
        model = JsonModel()
        model.json = {"check": 12}
        self.assertEqual(model.json, {"check": 12})
        self.assertEqual(type(model.json), dict)

        self.assertEqual(model.default_json, {"check": 12})
        self.assertEqual(type(model.default_json), dict)

    def test_invalid_json(self):
        # invalid json data {] in the json and default_json fields
        ser = '[{"pk": 1, "model": "jsonfield.jsoncharmodel", ' \
            '"fields": {"json": "{]", "default_json": "{]"}}]'
        with self.assertRaises(DeserializationError) as cm:
            next(deserialize('json', ser))
        inner = cm.exception.args[0]
        self.assertTrue(isinstance(inner, ValidationError))
        self.assertEqual('Enter valid JSON', inner.messages[0])

    def test_integer_in_string_in_json_field(self):
        """Test saving the Python string '123' in our JSONField"""
        json_obj = '123'
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_boolean_in_string_in_json_field(self):
        """Test saving the Python string 'true' in our JSONField"""
        json_obj = 'true'
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_pass_by_reference_pollution(self):
        """Make sure the default parameter is copied rather than passed by reference"""
        model = JsonModel()
        model.default_json["check"] = 144
        model.complex_default_json[0]["checkcheck"] = 144
        self.assertEqual(model.default_json["check"], 144)
        self.assertEqual(model.complex_default_json[0]["checkcheck"], 144)

        # Make sure when we create a new model, it resets to the default value
        # and not to what we just set it to (it would be if it were passed by reference)
        model = JsonModel()
        self.assertEqual(model.default_json["check"], 12)
        self.assertEqual(model.complex_default_json[0]["checkcheck"], 1212)

    def test_normal_regex_filter(self):
        """Make sure JSON model can filter regex"""

        JsonModel.objects.create(json={"boom": "town"})
        JsonModel.objects.create(json={"move": "town"})
        JsonModel.objects.create(json={"save": "town"})

        self.assertEqual(JsonModel.objects.count(), 3)

        self.assertEqual(JsonModel.objects.filter(json__regex=r"boom").count(), 1)
        self.assertEqual(JsonModel.objects.filter(json__regex=r"town").count(), 3)

class JSONCharFieldTest(JSONFieldTest):
    json_model = JsonCharModel

class OrderedJsonModel(models.Model):
    json = JSONField(load_kwargs={'object_pairs_hook': OrderedDict})

class OrderedDictSerializationTest(TestCase):
    ordered_dict = OrderedDict([
        ('number', [1, 2, 3, 4]),
        ('notes', True),
    ])
    expected_key_order = ['number', 'notes']

    def test_ordered_dict_differs_from_normal_dict(self):
        self.assertEqual(list(self.ordered_dict.keys()), self.expected_key_order)
        self.assertNotEqual(dict(self.ordered_dict).keys(), self.expected_key_order)

    def test_default_behaviour_loses_sort_order(self):
        mod = JsonModel.objects.create(json=self.ordered_dict)
        self.assertEqual(list(mod.json.keys()), self.expected_key_order)
        mod_from_db = JsonModel.objects.get(id=mod.id)

        # mod_from_db lost ordering information during json.loads()
        self.assertNotEqual(mod_from_db.json.keys(), self.expected_key_order)

    def test_load_kwargs_hook_does_not_lose_sort_order(self):
        mod = OrderedJsonModel.objects.create(json=self.ordered_dict)
        self.assertEqual(list(mod.json.keys()), self.expected_key_order)
        mod_from_db = OrderedJsonModel.objects.get(id=mod.id)
        self.assertEqual(list(mod_from_db.json.keys()), self.expected_key_order)
<commit_msg>Added new test for saving blank objects<commit_after>from decimal import Decimal
from django.core.serializers import deserialize, serialize
from django.core.serializers.base import DeserializationError
from django.db import models
from django.test import TestCase
from django.utils import simplejson as json

from .fields import JSONField, JSONCharField
from django.forms.util import ValidationError

from collections import OrderedDict

class JsonModel(models.Model):
    json = JSONField()
    default_json = JSONField(default={"check":12})
    complex_default_json = JSONField(default=[{"checkcheck": 1212}])
    empty_default = JSONField(default={})

class JsonCharModel(models.Model):
    json = JSONCharField(max_length=100)
    default_json = JSONCharField(max_length=100, default={"check":34})

class ComplexEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, complex):
            return {
                '__complex__': True,
                'real': obj.real,
                'imag': obj.imag,
            }

        return json.JSONEncoder.default(self, obj)

def as_complex(dct):
    if '__complex__' in dct:
        return complex(dct['real'], dct['imag'])
    return dct

class JSONModelCustomEncoders(models.Model):
    # A JSON field that can store complex numbers
    json = JSONField(
        dump_kwargs={'cls': ComplexEncoder, "indent": 4},
        load_kwargs={'object_hook': as_complex},
    )

class JSONFieldTest(TestCase):
    """JSONField Wrapper Tests"""

    json_model = JsonModel

    def test_json_field_create(self):
        """Test saving a JSON object in our JSONField"""
        json_obj = {
            "item_1": "this is a json blah",
            "blergh": "hey, hey, hey"}

        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_string_in_json_field(self):
        """Test saving an ordinary Python string in our JSONField"""
        json_obj = 'blah blah'
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_float_in_json_field(self):
        """Test saving a Python float in our JSONField"""
        json_obj = 1.23
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_int_in_json_field(self):
        """Test saving a Python integer in our JSONField"""
        json_obj = 1234567
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_decimal_in_json_field(self):
        """Test saving a Python Decimal in our JSONField"""
        json_obj = Decimal(12.34)
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        # here we must know to convert the returned string back to Decimal,
        # since json does not support that format
        self.assertEqual(Decimal(new_obj.json), json_obj)

    def test_json_field_modify(self):
        """Test modifying a JSON object in our JSONField"""
        json_obj_1 = {'a': 1, 'b': 2}
        json_obj_2 = {'a': 3, 'b': 4}

        obj = self.json_model.objects.create(json=json_obj_1)
        self.assertEqual(obj.json, json_obj_1)
        obj.json = json_obj_2

        self.assertEqual(obj.json, json_obj_2)
        obj.save()
        self.assertEqual(obj.json, json_obj_2)

        self.assertTrue(obj)

    def test_json_field_load(self):
        """Test loading a JSON object from the DB"""
        json_obj_1 = {'a': 1, 'b': 2}
        obj = self.json_model.objects.create(json=json_obj_1)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj_1)

    def test_json_list(self):
        """Test storing a JSON list"""
        json_obj = ["my", "list", "of", 1, "objs", {"hello": "there"}]

        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)
        self.assertEqual(new_obj.json, json_obj)

    def test_empty_objects(self):
        """Test storing empty objects"""
        for json_obj in [{}, [], 0, '', False]:
            obj = self.json_model.objects.create(json=json_obj)
            new_obj = self.json_model.objects.get(id=obj.id)
            self.assertEqual(json_obj, obj.json)
            self.assertEqual(json_obj, new_obj.json)

    def test_custom_encoder(self):
        """Test encoder_cls and object_hook"""
        value = 1 + 3j  # A complex number

        obj = JSONModelCustomEncoders.objects.create(json=value)
        new_obj = JSONModelCustomEncoders.objects.get(pk=obj.pk)
        self.assertEqual(value, new_obj.json)

    def test_django_serializers(self):
        """Test serializing/deserializing jsonfield data"""
        for json_obj in [{}, [], 0, '', False, {'key': 'value', 'num': 42,
                                                'ary': list(range(5)),
                                                'dict': {'k': 'v'}}]:
            obj = self.json_model.objects.create(json=json_obj)
            new_obj = self.json_model.objects.get(id=obj.id)
            self.assert_(new_obj)

        queryset = self.json_model.objects.all()
        ser = serialize('json', queryset)
        for dobj in deserialize('json', ser):
            obj = dobj.object
            pulled = self.json_model.objects.get(id=obj.pk)
            self.assertEqual(obj.json, pulled.json)

    def test_default_parameters(self):
        """Test providing a default value to the model"""
        model = JsonModel()
        model.json = {"check": 12}
        self.assertEqual(model.json, {"check": 12})
        self.assertEqual(type(model.json), dict)

        self.assertEqual(model.default_json, {"check": 12})
        self.assertEqual(type(model.default_json), dict)

    def test_invalid_json(self):
        # invalid json data {] in the json and default_json fields
        ser = '[{"pk": 1, "model": "jsonfield.jsoncharmodel", ' \
            '"fields": {"json": "{]", "default_json": "{]"}}]'
        with self.assertRaises(DeserializationError) as cm:
            next(deserialize('json', ser))
        inner = cm.exception.args[0]
        self.assertTrue(isinstance(inner, ValidationError))
        self.assertEqual('Enter valid JSON', inner.messages[0])

    def test_integer_in_string_in_json_field(self):
        """Test saving the Python string '123' in our JSONField"""
        json_obj = '123'
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_boolean_in_string_in_json_field(self):
        """Test saving the Python string 'true' in our JSONField"""
        json_obj = 'true'
        obj = self.json_model.objects.create(json=json_obj)
        new_obj = self.json_model.objects.get(id=obj.id)

        self.assertEqual(new_obj.json, json_obj)

    def test_pass_by_reference_pollution(self):
        """Make sure the default parameter is copied rather than passed by reference"""
        model = JsonModel()
        model.default_json["check"] = 144
        model.complex_default_json[0]["checkcheck"] = 144
        self.assertEqual(model.default_json["check"], 144)
        self.assertEqual(model.complex_default_json[0]["checkcheck"], 144)

        # Make sure when we create a new model, it resets to the default value
        # and not to what we just set it to (it would be if it were passed by reference)
        model = JsonModel()
        self.assertEqual(model.default_json["check"], 12)
        self.assertEqual(model.complex_default_json[0]["checkcheck"], 1212)

    def test_normal_regex_filter(self):
        """Make sure JSON model can filter regex"""

        JsonModel.objects.create(json={"boom": "town"})
        JsonModel.objects.create(json={"move": "town"})
        JsonModel.objects.create(json={"save": "town"})

        self.assertEqual(JsonModel.objects.count(), 3)

        self.assertEqual(JsonModel.objects.filter(json__regex=r"boom").count(), 1)
        self.assertEqual(JsonModel.objects.filter(json__regex=r"town").count(), 3)

    def test_save_blank_object(self):
        """Test that JSON model can save a blank object as none"""

        model = JsonModel()
        self.assertEqual(model.empty_default, {})

        model.save()
        self.assertEqual(model.empty_default, {})

        model1 = JsonModel(empty_default={"hey": "now"})
        self.assertEqual(model1.empty_default, {"hey": "now"})

        model1.save()
        self.assertEqual(model1.empty_default, {"hey": "now"})

class JSONCharFieldTest(JSONFieldTest):
    json_model = JsonCharModel

class OrderedJsonModel(models.Model):
    json = JSONField(load_kwargs={'object_pairs_hook': OrderedDict})

class OrderedDictSerializationTest(TestCase):
    ordered_dict = OrderedDict([
        ('number', [1, 2, 3, 4]),
        ('notes', True),
    ])
    expected_key_order = ['number', 'notes']

    def test_ordered_dict_differs_from_normal_dict(self):
        self.assertEqual(list(self.ordered_dict.keys()), self.expected_key_order)
        self.assertNotEqual(dict(self.ordered_dict).keys(), self.expected_key_order)

    def test_default_behaviour_loses_sort_order(self):
        mod = JsonModel.objects.create(json=self.ordered_dict)
        self.assertEqual(list(mod.json.keys()), self.expected_key_order)
        mod_from_db = JsonModel.objects.get(id=mod.id)

        # mod_from_db lost ordering information during json.loads()
        self.assertNotEqual(mod_from_db.json.keys(), self.expected_key_order)

    def test_load_kwargs_hook_does_not_lose_sort_order(self):
        mod = OrderedJsonModel.objects.create(json=self.ordered_dict)
        self.assertEqual(list(mod.json.keys()), self.expected_key_order)
        mod_from_db = OrderedJsonModel.objects.get(id=mod.id)
        self.assertEqual(list(mod_from_db.json.keys()), self.expected_key_order)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# XDS2CBF.py
#
#   Copyright (C) 2010 Diamond Light Source, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# Code to read in a GXPARM file from XDS and write out the corresponding
# CBF reference frame all nicely refined.

from scitbx import matrix
import math
import sys
import os

if not os.environ.has_key('XIA2_ROOT'):
    raise RuntimeError, 'XIA2_ROOT not defined'

if not os.path.join(os.environ['XIA2_ROOT']) in sys.path:
    sys.path.append(os.path.join(os.environ['XIA2_ROOT']))

from Wrappers.XDS.XDS import xds_read_xparm

def parse_xparm(xparm_file):
    '''Read an xparm file, return the rotation axis and beam vector in the
    XDS coordinate frame. Also the vector in this same coordinate frame to
    the start of the first pixel on the detector.'''

    xdata = xds_read_xparm(xparm_file)

    ra = xdata['axis']
    beam = xdata['beam']
    distance = xdata['distance']
    px = xdata['px']
    py = xdata['py']
    ox = xdata['ox']
    oy = xdata['oy']
    nx = xdata['nx']
    ny = xdata['ny']
    x_to_d = - px * ox, - py * oy, - distance

    x = xdata['x']
    y = xdata['y']

    x_to_d = - matrix.col(x) * px * ox + \
             - matrix.col(y) * py * oy + \
             distance * matrix.col(x).cross(matrix.col(y))

    return ra, beam, x_to_d.elems, (px, py), distance, (nx, ny), x, y

def get_abc_from_xparm(xparm_file):

    xdata = xds_read_xparm(xparm_file)

    return xdata['a'], xdata['b'], xdata['c']

def XDS2CBF(xparm_file, integrate_hkl):
    '''Given an XDS XPARM file, return a matrix which will transform from
    XDS coordinates to CBF reference frame.'''

    # XDS coordinate frame: x = detector fast direction
    #                       y = detector slow direction
    #                       z = x ^ y is detector normal (n.b. TOWARDS
    #                                                     the source)
    # CBF coordinate frame: x = primary gonio axis
    #                       y = z ^ x
    #                       z = component of beam |'r to x
    #
    # well that's ok then...

    ra, beam, x_to_d, pixel, distance, nxny, dx, dy = parse_xparm(xparm_file)

    # make them vectors etc.

    ra = matrix.col(ra)
    beam = matrix.col(beam)
    x_to_d = matrix.col(x_to_d)
    nx, ny = nxny

    wavelength = 1.0 / math.sqrt(beam.dot())

    # then unit vectors
    ra = ra / math.sqrt(ra.dot())
    beam = beam / math.sqrt(beam.dot())

    # right now lets calculate a rotation which will overlap the two reference
    # frames.... first relabel so we know what we are trying to do

    x = ra
    z = -1 * (beam - (beam.dot(ra) * ra))
    z = z / math.sqrt(z.dot())
    y = z.cross(x)

    # now lets figure the rotations we want to rotate x to _x, etc.

    _x = matrix.col([1, 0, 0])
    _y = matrix.col([0, 1, 0])
    _z = matrix.col([0, 0, 1])

    # ok then - #1 rotate about x ^ (1, 0, 0) - if they are not identical

    if _x.dot(x):
        _ra_x = _x.cross(x)
        _a_x = _x.angle(x)
    else:
        _ra_x = _x
        _a_x = 0.0

    _m_x = _ra_x.axis_and_angle_as_r3_rotation_matrix(- _a_x)

    # then rotate z to _z by rotating about _x (which is now coincident
    # with x) N.B. z cannot be perpendicular to x

    _ra_z = _x
    _a_z = _z.angle(_m_x * z)

    _m_z = _ra_z.axis_and_angle_as_r3_rotation_matrix(- _a_z)

    # _m is matrix to rotate FROM xds coordinate frame TO cbf reference frame

    _m = _m_z * _m_x

    # verify that this is a rotation i.e. has determinant +1.

    assert(math.fabs(_m.determinant() - 1.0) < 1.0e-7)

    # and finally the new beam

    print 'New beam vector'

    print '%10.7f %10.7f %10.7f' % (_m * beam).elems

    print 'To detector origin i.e. start of the first pixel'

    print '%10.4f %10.4f %10.4f' % (_m * x_to_d).elems

    # now need to consider the position of the detector etc. to derive the
    # new direct beam centre...

    _X = _m * dx
    _Y = _m * dy
    _N = _X.cross(_Y)

    # so that would be the distance between the sample and the
    # detector plane in the direction of the direct beam.

    _beam = _m * beam

    # D vector to detector origin, B vector to intersection of beam with det

    D = _m * x_to_d
    d = _N.dot(D)

    print 'Distance is %f' % d

    B = _beam * (d / _beam.dot(_N))

    o = B - D

    beam_x = o.dot(_X)
    beam_y = o.dot(_Y)
    beam_n = o.dot(_N)

    print 'Beam centre on the detector'
    print '%10.4f %10.4f %10.4f' % (beam_x, beam_y, beam_n)

    # check the result is in the detector plane
    assert(math.fabs(beam_n) < 1.0e-7)

    print 'Detector axes'
    print '%10.4f %10.4f %10.4f' % _X.elems
    print '%10.4f %10.4f %10.4f' % _Y.elems

    a, b, c = get_abc_from_xparm(xparm_file)

    _a = _m * a
    _b = _m * b
    _c = _m * c

    UB = matrix.sqr(_a.elems + _b.elems + _c.elems).inverse()

    start_angle = None
    angle_range = None
    start_frame = None

    S0 = (1 / wavelength) * _beam
    O = _m * x_to_d

    for record in open(integrate_hkl):
        if '!' in record[:1]:
            if '!STARTING_ANGLE' in record:
                start_angle = float(record.split()[-1])
            elif '!STARTING_FRAME' in record:
                start_frame = float(record.split()[-1])
            elif '!OSCILLATION_RANGE' in record:
                angle_range = float(record.split()[-1])
            continue

        hkl = tuple(map(int, record.split()[:3]))
        xyz = tuple(map(float, record.split()[5:8]))

        phi = (xyz[2] - start_frame) * angle_range + start_angle

        R = _x.axis_and_angle_as_r3_rotation_matrix(math.pi * phi / 180.0)

        q = R * UB * hkl

        p = S0 + q

        p_ = p * (1.0 / math.sqrt(p.dot()))
        P = p_ * (O.dot(_N) / (p_.dot(_N)))

        R = P - O

        i = R.dot(_X)
        j = R.dot(_Y)
        k = R.dot(_N)

        print '%d %d %d %.3f %.3f %.3f %.3f %.3f %.3f' % \
              (hkl[0], hkl[1], hkl[2], i, j, k,
               xyz[0] * pixel[0], xyz[1] * pixel[1], phi)

if __name__ == '__main__':
    XDS2CBF(sys.argv[1], sys.argv[2])
<commit_msg>nice, now works<commit_after>#!/usr/bin/env python
# XDS2CBF.py
#
#   Copyright (C) 2010 Diamond Light Source, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# Code to read in a GXPARM file from XDS and write out the corresponding
# CBF reference frame all nicely refined.

from scitbx import matrix
import math
import sys
import os

if not os.environ.has_key('XIA2_ROOT'):
    raise RuntimeError, 'XIA2_ROOT not defined'

if not os.path.join(os.environ['XIA2_ROOT']) in sys.path:
    sys.path.append(os.path.join(os.environ['XIA2_ROOT']))

from Wrappers.XDS.XDS import xds_read_xparm

def parse_xparm(xparm_file):
    '''Read an xparm file, return the rotation axis and beam vector in the
    XDS coordinate frame. Also the vector in this same coordinate frame to
    the start of the first pixel on the detector.'''

    xdata = xds_read_xparm(xparm_file)

    ra = xdata['axis']
    beam = xdata['beam']
    distance = xdata['distance']
    px = xdata['px']
    py = xdata['py']
    ox = xdata['ox']
    oy = xdata['oy']
    nx = xdata['nx']
    ny = xdata['ny']
    x_to_d = - px * ox, - py * oy, - distance

    x = xdata['x']
    y = xdata['y']

    x_to_d = - matrix.col(x) * px * ox + \
             - matrix.col(y) * py * oy + \
             distance * matrix.col(x).cross(matrix.col(y))

    return ra, beam, x_to_d.elems, (px, py), distance, (nx, ny), x, y

def get_abc_from_xparm(xparm_file):

    xdata = xds_read_xparm(xparm_file)

    return xdata['a'], xdata['b'], xdata['c']

def XDS2CBF(xparm_file, integrate_hkl):
    '''Given an XDS XPARM file, return a matrix which will transform from
    XDS coordinates to CBF reference frame.'''

    # XDS coordinate frame: x = detector fast direction
    #                       y = detector slow direction
    #                       z = x ^ y is detector normal (n.b. TOWARDS
    #                                                     the source)
    # CBF coordinate frame: x = primary gonio axis
    #                       y = z ^ x
    #                       z = component of beam |'r to x
    #
    # well that's ok then...

    ra, beam, x_to_d, pixel, distance, nxny, dx, dy = parse_xparm(xparm_file)

    # make them vectors etc.

    ra = matrix.col(ra)
    beam = matrix.col(beam)
    x_to_d = matrix.col(x_to_d)
    nx, ny = nxny

    wavelength = 1.0 / math.sqrt(beam.dot())

    # then unit vectors
    ra = ra / math.sqrt(ra.dot())
    beam = beam / math.sqrt(beam.dot())

    # right now lets calculate a rotation which will overlap the two reference
    # frames.... first relabel so we know what we are trying to do

    x = ra
    z = -1 * (beam - (beam.dot(ra) * ra))
    z = z / math.sqrt(z.dot())
    y = z.cross(x)

    # now lets figure the rotations we want to rotate x to _x, etc.

    _x = matrix.col([1, 0, 0])
    _y = matrix.col([0, 1, 0])
    _z = matrix.col([0, 0, 1])

    # ok then - #1 rotate about x ^ (1, 0, 0) - if they are not identical

    if _x.dot(x):
        _ra_x = _x.cross(x)
        _a_x = _x.angle(x)
    else:
        _ra_x = _x
        _a_x = 0.0

    _m_x = _ra_x.axis_and_angle_as_r3_rotation_matrix(- _a_x)

    # then rotate z to _z by rotating about _x (which is now coincident
    # with x) N.B. z cannot be perpendicular to x

    _ra_z = _x
    _a_z = _z.angle(_m_x * z)

    _m_z = _ra_z.axis_and_angle_as_r3_rotation_matrix(- _a_z)

    # _m is matrix to rotate FROM xds coordinate frame TO cbf reference frame

    _m = _m_z * _m_x

    # verify that this is a rotation i.e. has determinant +1.

    assert(math.fabs(_m.determinant() - 1.0) < 1.0e-7)

    # and finally the new beam

    print 'New beam vector'

    print '%10.7f %10.7f %10.7f' % (_m * beam).elems

    print 'To detector origin i.e. start of the first pixel'

    print '%10.4f %10.4f %10.4f' % (_m * x_to_d).elems

    # now need to consider the position of the detector etc. to derive the
    # new direct beam centre...

    _X = _m * dx
    _Y = _m * dy
    _N = _X.cross(_Y)

    # so that would be the distance between the sample and the
    # detector plane in the direction of the direct beam.

    _beam = _m * beam

    # D vector to detector origin, B vector to intersection of beam with det

    D = _m * x_to_d
    d = _N.dot(D)

    print 'Distance is %f' % d

    B = _beam * (d / _beam.dot(_N))

    o = B - D

    beam_x = o.dot(_X)
    beam_y = o.dot(_Y)
    beam_n = o.dot(_N)

    print 'Beam centre on the detector'
    print '%10.4f %10.4f %10.4f' % (beam_x, beam_y, beam_n)

    # check the result is in the detector plane
    assert(math.fabs(beam_n) < 1.0e-7)

    print 'Detector axes'
    print '%10.4f %10.4f %10.4f' % _X.elems
    print '%10.4f %10.4f %10.4f' % _Y.elems

    a, b, c = get_abc_from_xparm(xparm_file)

    _a = _m * a
    _b = _m * b
    _c = _m * c

    UB = matrix.sqr(_a.elems + _b.elems + _c.elems).inverse()

    print 'UB matrix:'
    print '%10.7f %10.7f %10.7f' % UB.elems[0:3]
    print '%10.7f %10.7f %10.7f' % UB.elems[3:6]
    print '%10.7f %10.7f %10.7f' % UB.elems[6:9]

    start_angle = None
    angle_range = None
    start_frame = None

    S0 = (1 / wavelength) * _beam
    O = _m * x_to_d

    for record in open(integrate_hkl):
        if '!' in record[:1]:
            if '!STARTING_ANGLE' in record:
                start_angle = float(record.split()[-1])
            elif '!STARTING_FRAME' in record:
                start_frame = float(record.split()[-1])
            elif '!OSCILLATION_RANGE' in record:
                angle_range = float(record.split()[-1])
            continue

        hkl = tuple(map(int, record.split()[:3]))
        xyz = tuple(map(float, record.split()[5:8]))

        phi = (xyz[2] - start_frame) * angle_range + start_angle

        R = _x.axis_and_angle_as_r3_rotation_matrix(math.pi * phi / 180.0)

        q = R * UB * hkl

        p = S0 + q

        p_ = p * (1.0 / math.sqrt(p.dot()))
        P = p_ * (O.dot(_N) / (p_.dot(_N)))

        R = P - O

        i = R.dot(_X)
        j = R.dot(_Y)
        k = R.dot(_N)

        print '%d %d %d %.3f %.3f %.3f %.3f %.3f %.3f' % \
              (hkl[0], hkl[1], hkl[2], i, j, k,
               xyz[0] * pixel[0], xyz[1] * pixel[1], phi)

        break

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


<commit_before># Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

import getpass
import os
import StringIO
import subprocess
import sys
import unittest

from django.core import exceptions

from openquake.engine import engine2
from openquake.engine.db import models
from openquake.engine.job import validation

from tests.utils import helpers

class PrepareJobTestCase(unittest.TestCase):

    def test_prepare_job_default_user(self):
        job = engine2.prepare_job()

        self.assertEqual('openquake', job.owner.user_name)
        self.assertEqual('pre_executing', job.status)
        self.assertEqual('progress', job.log_level)

        # Check the make sure it's in the database.
        try:
            models.OqJob.objects.get(id=job.id)
        except exceptions.ObjectDoesNotExist:
            self.fail('Job was not found in the database')

    def test_prepare_job_specified_user(self):
        user_name = helpers.random_string()
        job = engine2.prepare_job(user_name=user_name)

        self.assertEqual(user_name, job.owner.user_name)
        self.assertEqual('pre_executing', job.status)
        self.assertEqual('progress', job.log_level)

        try:
            models.OqJob.objects.get(id=job.id)
        except exceptions.ObjectDoesNotExist:
            self.fail('Job was not found in the database')

    def test_prepare_job_explicit_log_level(self):
        # By default, a job is created with a log level of 'progress'
        # (just to show calculation progress).
        # In this test, we'll specify 'debug' as the log level.
        job = engine2.prepare_job(log_level='debug')

        self.assertEqual('debug', job.log_level)

class PrepareUserTestCase(unittest.TestCase):

    def test_prepare_user_exists(self):
        user_name = helpers.random_string()
        existing_user = models.OqUser(
            user_name=user_name, full_name=user_name, organization_id=1
        )
        existing_user.save()

        user = engine2.prepare_user(user_name)
        self.assertEqual(existing_user.id, user.id)

    def test_prepare_user_does_not_exist(self):
        user_name = helpers.random_string()

        # Sanity check; make sure the user doesn't exist yet.
        self.assertEqual(
            0, len(models.OqUser.objects.filter(user_name=user_name))
        )

        engine2.prepare_user(user_name)

        # Now the user should exist.
        self.assertEqual(
            1, len(models.OqUser.objects.filter(user_name=user_name))
        )

class ParseConfigTestCase(unittest.TestCase):

    def test_parse_config_no_files(self):
        # sections are there just for documentation
        # when we parse the file, we ignore these
        source = StringIO.StringIO("""
[general]
CALCULATION_MODE = classical
region = 1 1 2 2 3 3
[foo]
bar = baz
""")

        # Add a 'name' to make this look like a real file:
        source.name = 'path/to/some/job.ini'
        exp_base_path = os.path.dirname(
            os.path.join(os.path.abspath('.'), source.name))

        expected_params = {
            'base_path': exp_base_path,
            'force_inputs': False,
            'calculation_mode': 'classical',
            'region': '1 1 2 2 3 3',
            'bar': 'baz',
        }

        params, _ = engine2.parse_config(source)

        self.assertEqual(expected_params, params)

    def test_parse_config_with_files_force_inputs(self):
        site_model_input = helpers.touch(content="site model")
        sm_lt_input = helpers.touch(content="source model logic tree")
        gsim_lt_input = helpers.touch(content="gsim logic tree")

        source = StringIO.StringIO("""
[hazard_or_whatever]
calculation_mode = classical
gsim_logic_tree_file = %s
source_model_logic_tree_file = %s
site_model_file = %s
not_a_valid_xml = foo.xml
""" % (gsim_lt_input, sm_lt_input, site_model_input))

        # Add a 'name' to make this look like a real file:
        source.name = 'path/to/some/job.ini'
        exp_base_path = os.path.dirname(
            os.path.join(os.path.abspath('.'), source.name))

        expected_params = {
            'base_path': exp_base_path,
            'force_inputs': True,
            'calculation_mode': 'classical',
            'not_a_valid_xml': 'foo.xml',
        }

        params, files = engine2.parse_config(source, force_inputs=True)

        expected_files = {
            'site_model_file': models.Input.objects.filter(
                input_type='site_model').latest('id'),
            'source_model_logic_tree_file': models.Input.objects.filter(
                input_type='source_model_logic_tree').latest('id'),
            'gsim_logic_tree_file': models.Input.objects.filter(
                input_type='gsim_logic_tree').latest('id'),
        }

        self.assertEqual(expected_params, params)

        self.assertEqual(len(expected_files), len(files))
        for key in expected_files:
            self.assertEqual(expected_files[key].id, files[key].id)

    def test_parse_config_with_files_no_force_inputs(self):
        site_model_input = helpers.touch(content="foo")

        source = StringIO.StringIO("""
[general]
calculation_mode = classical
[site]
site_model_file = %s
maximum_distance=0
truncation_level=0
random_seed=0
""" % site_model_input)

        # Add a 'name' to make this look like a real file:
        source.name = 'path/to/some/job.ini'
        exp_base_path = os.path.dirname(
            os.path.join(os.path.abspath('.'), source.name))

        expected_params = {
            'base_path': exp_base_path,
            'force_inputs': False,
            'calculation_mode': 'classical',
            'truncation_level': '0',
            'random_seed': '0',
            'maximum_distance': '0'
        }

        # Run first with force_inputs=True to create the new Input.
        params, expected_files = engine2.parse_config(
            source, force_inputs=True)

        # In order for us to reuse the existing input, we need to associate
        # each input with a successful job.
        job = engine2.prepare_job(getpass.getuser())

        job.hazard_calculation = engine2.create_hazard_calculation(
            job.owner, params, expected_files.values())
        job.status = 'complete'
        job.save()

        # Now run twice with force_inputs=False (the default).
        source.seek(0)
        params1, files1 = engine2.parse_config(source)
        source.seek(0)
        params2, files2 = engine2.parse_config(source)

        # Check the params just for sanity.
        self.assertEqual(expected_params, params1)
        self.assertEqual(expected_params, params2)

        # Finally, check that the Input returned by the latest 2 calls matches
        # the input we created above.
        self.assertEqual(len(expected_files), len(files1))
        self.assertEqual(len(expected_files), len(files2))

        self.assertEqual(
            expected_files['site_model_file'].id, files1['site_model_file'].id)
        self.assertEqual(
            expected_files['site_model_file'].id, files2['site_model_file'].id)

class FileDigestTestCase(unittest.TestCase):
    """Test the _file_digest() function."""

    PATH = helpers.get_data_path("src_model1.dat")

    def test__file_digest(self):
        # Make sure the digest returned by the function matches the one
        # obtained via /usr/bin/md5sum
        if sys.platform == 'darwin':
            expected = subprocess.check_output(["md5", self.PATH]).split()[-1]
        else:
            expected = subprocess.check_output(
                ["md5sum", self.PATH]).split()[0]
        actual = engine2._file_digest(self.PATH)
        self.assertEqual(expected, actual)

class GetContentTypeTestCase(unittest.TestCase):

    def test__get_content_type(self):
        # no file extension
        self.assertEqual('unknown', engine2._get_content_type('/foo/bar/baz'))
        # xml file extension
        self.assertEqual('xml', engine2._get_content_type('/foo/bar/baz.xml'))
        # hdf5 file extension
        self.assertEqual(
            'HDF5', engine2._get_content_type('/foo/bar/baz.HDF5')
        )

class CreateHazardCalculationTestCase(unittest.TestCase):

    def test_create_hazard_calculation(self):
        # Just the bare minimum set of params to satisfy not null constraints
        # in the db.
        params = {
            'base_path': 'path/to/job.ini',
            'calculation_mode': 'classical',
            'region': '1 1 2 2 3 3',
            'width_of_mfd_bin': '1',
            'rupture_mesh_spacing': '1',
            'area_source_discretization': '2',
            'investigation_time': 50,
            'truncation_level': 0,
            'maximum_distance': 200,
            'number_of_logic_tree_samples': 1,
            'intensity_measure_types_and_levels': dict(PGA=[1, 2, 3, 4]),
            'random_seed': 37,
        }

        owner = helpers.default_user()

        site_model = models.Input(digest='123', path='/foo/bar', size=0,
                                  input_type='site_model', owner=owner)
        site_model.save()
        files = [site_model]

        hc = engine2.create_hazard_calculation(owner, params, files)
        # Normalize/clean fields by fetching a fresh copy from the db.
        hc = models.HazardCalculation.objects.get(id=hc.id)

        self.assertEqual(hc.calculation_mode, 'classical')
        self.assertEqual(hc.width_of_mfd_bin, 1.0)
        self.assertEqual(hc.rupture_mesh_spacing, 1.0)
        self.assertEqual(hc.area_source_discretization, 2.0)
        self.assertEqual(hc.investigation_time, 50.0)
        self.assertEqual(hc.truncation_level, 0.0)
        self.assertEqual(hc.maximum_distance, 200.0)

        # Test the input2haz_calc link:
        [inp2hcs] = models.Input2hcalc.objects.filter(
            hazard_calculation=hc.id)

        self.assertEqual(site_model.id, inp2hcs.input.id)

class CreateRiskCalculationTestCase(unittest.TestCase):

    def test_create_risk_calculation(self):
        params = {
            'hazard_output_id': 1,
            'base_path': 'path/to/job.ini',
            'export_dir': '/tmp/xxx',
            'calculation_mode': 'classical',
            # just some sample params
            'lrem_steps_per_interval': 5,
            'conditional_loss_poes': '0.01, 0.02, 0.05',
            'region_constraint': '-0.5 0.5, 0.5 0.5, 0.5 -0.5, -0.5, -0.5',
        }

        owner = helpers.default_user()

        vuln_file = models.Input(digest='123', path='/foo/bar', size=0,
                                 input_type='vulnerability', owner=owner)
        vuln_file.save()
        exposure_file = models.Input(digest='456', path='/foo/baz', size=0,
                                     input_type='exposure', owner=owner)
        exposure_file.save()

        files = [vuln_file, exposure_file]

        rc = engine2.create_risk_calculation(owner, params, files)
        # Normalize/clean fields by fetching a fresh copy from the db.
        rc = models.RiskCalculation.objects.get(id=rc.id)

        self.assertEqual(rc.calculation_mode, 'classical')
        self.assertEqual(rc.lrem_steps_per_interval, 5)
        self.assertEqual(rc.conditional_loss_poes, [0.01, 0.02, 0.05])
        self.assertEqual(
            rc.region_constraint.wkt,
            ('POLYGON ((-0.5000000000000000 0.5000000000000000, '
             '0.5000000000000000 0.5000000000000000, '
             '0.5000000000000000 -0.5000000000000000, '
             '-0.5000000000000000 -0.5000000000000000, '
             '-0.5000000000000000 0.5000000000000000))'))

class ReadJobProfileFromConfigFileTestCase(unittest.TestCase):
    """Integration test for basic engine functions.

    Test reading/generating a hazard job profile from a config file, then run
    through the validation form.
    """

    def test_read_and_validate_hazard_config(self):
        cfg = helpers.demo_file('simple_fault_demo_hazard/job.ini')
        job = engine2.prepare_job(getpass.getuser())
        params, files = engine2.parse_config(open(cfg, 'r'))
        calculation = engine2.create_hazard_calculation(
            job.owner, params, files.values())

        form = validation.ClassicalHazardForm(
            instance=calculation, files=files
        )
        self.assertTrue(form.is_valid())

class OpenquakeCliTestCase(unittest.TestCase):
    """
    Run "openquake --version" as a separate
    process using `subprocess`.
    """

    def test_run_version(self):
        args = [helpers.RUNNER, "--version"]

        print 'Running:', ' '.join(args)  # this is useful for debugging
        return subprocess.check_call(args)
<commit_msg>tests/engine2_test:<commit_after># Copyright (c) 2010-2012, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

import getpass
import os
import StringIO
import subprocess
import sys
import unittest

from django.core import exceptions

from openquake.engine import engine2
from openquake.engine.db import models
from openquake.engine.job import validation

from tests.utils import helpers

class PrepareJobTestCase(unittest.TestCase):

    def test_prepare_job_default_user(self):
        job = engine2.prepare_job()

        self.assertEqual('openquake', job.owner.user_name)
        self.assertEqual('pre_executing', job.status)
        self.assertEqual('progress', job.log_level)

        # Check the make sure it's in the database.
        try:
            models.OqJob.objects.get(id=job.id)
        except exceptions.ObjectDoesNotExist:
            self.fail('Job was not found in the database')

    def test_prepare_job_specified_user(self):
        user_name = helpers.random_string()
        job = engine2.prepare_job(user_name=user_name)

        self.assertEqual(user_name, job.owner.user_name)
        self.assertEqual('pre_executing', job.status)
        self.assertEqual('progress', job.log_level)

        try:
            models.OqJob.objects.get(id=job.id)
        except exceptions.ObjectDoesNotExist:
            self.fail('Job was not found in the database')

    def test_prepare_job_explicit_log_level(self):
        # By default, a job is created with a log level of 'progress'
        # (just to show calculation progress).
        # In this test, we'll specify 'debug' as the log level.
        job = engine2.prepare_job(log_level='debug')

        self.assertEqual('debug', job.log_level)

class PrepareUserTestCase(unittest.TestCase):

    def test_prepare_user_exists(self):
        user_name = helpers.random_string()
        existing_user = models.OqUser(
            user_name=user_name, full_name=user_name, organization_id=1
        )
        existing_user.save()

        user = engine2.prepare_user(user_name)
        self.assertEqual(existing_user.id, user.id)

    def test_prepare_user_does_not_exist(self):
        user_name = helpers.random_string()

        # Sanity check; make sure the user doesn't exist yet.
        self.assertEqual(
            0, len(models.OqUser.objects.filter(user_name=user_name))
        )

        engine2.prepare_user(user_name)

        # Now the user should exist.
        self.assertEqual(
            1, len(models.OqUser.objects.filter(user_name=user_name))
        )

class ParseConfigTestCase(unittest.TestCase):

    def test_parse_config_no_files(self):
        # sections are there just for documentation
        # when we parse the file, we ignore these
        source = StringIO.StringIO("""
[general]
CALCULATION_MODE = classical
region = 1 1 2 2 3 3
[foo]
bar = baz
""")

        # Add a 'name' to make this look like a real file:
        source.name = 'path/to/some/job.ini'
        exp_base_path = os.path.dirname(
            os.path.join(os.path.abspath('.'), source.name))

        expected_params = {
            'base_path': exp_base_path,
            'force_inputs': False,
            'calculation_mode': 'classical',
            'region': '1 1 2 2 3 3',
            'bar': 'baz',
        }

        params, _ = engine2.parse_config(source)

        self.assertEqual(expected_params, params)

    def test_parse_config_with_files_force_inputs(self):
        site_model_input = helpers.touch(content="site model")
        sm_lt_input = helpers.touch(content="source model logic tree")
        gsim_lt_input = helpers.touch(content="gsim logic tree")

        source = StringIO.StringIO("""
[hazard_or_whatever]
calculation_mode = classical
gsim_logic_tree_file = %s
source_model_logic_tree_file = %s
site_model_file = %s
not_a_valid_xml = foo.xml
""" % (gsim_lt_input, sm_lt_input, site_model_input))

        # Add a 'name' to make this look like a real file:
        source.name = 'path/to/some/job.ini'
        exp_base_path = os.path.dirname(
            os.path.join(os.path.abspath('.'), source.name))

        expected_params = {
            'base_path': exp_base_path,
            'force_inputs': True,
            'calculation_mode': 'classical',
            'not_a_valid_xml': 'foo.xml',
        }

        params, files = engine2.parse_config(source, force_inputs=True)

        expected_files = {
            'site_model_file': models.Input.objects.filter(
                input_type='site_model').latest('id'),
            'source_model_logic_tree_file': models.Input.objects.filter(
                input_type='source_model_logic_tree').latest('id'),
            'gsim_logic_tree_file': models.Input.objects.filter(
                input_type='gsim_logic_tree').latest('id'),
        }

        self.assertEqual(expected_params, params)

        self.assertEqual(len(expected_files), len(files))
        for key in expected_files:
            self.assertEqual(expected_files[key].id, files[key].id)

    def test_parse_config_with_files_no_force_inputs(self):
        site_model_input = helpers.touch(content="foo")

        source = StringIO.StringIO("""
[general]
calculation_mode = classical
[site]
site_model_file = %s
maximum_distance=0
truncation_level=0
random_seed=0
""" % site_model_input)

        # Add a 'name' to make this look like a real file:
        source.name = 'path/to/some/job.ini'
        exp_base_path = os.path.dirname(
            os.path.join(os.path.abspath('.'), source.name))

        expected_params = {
            'base_path': exp_base_path,
            'force_inputs': False,
            'calculation_mode': 'classical',
            'truncation_level': '0',
            'random_seed': '0',
            'maximum_distance': '0'
        }

        # Run first with force_inputs=True to create the new Input.
        params, expected_files = engine2.parse_config(
            source, force_inputs=True)

        # In order for us to reuse the existing input, we need to associate
        # each input with a successful job.
        job = engine2.prepare_job(getpass.getuser())

        job.hazard_calculation = engine2.create_hazard_calculation(
            job.owner, params, expected_files.values())
        job.status = 'complete'
        job.save()

        # Now run twice with force_inputs=False (the default).
        source.seek(0)
        params1, files1 = engine2.parse_config(source)
        source.seek(0)
        params2, files2 = engine2.parse_config(source)

        # Check the params just for sanity.
        self.assertEqual(expected_params, params1)
        self.assertEqual(expected_params, params2)

        # Finally, check that the Input returned by the latest 2 calls matches
        # the input we created above.
        self.assertEqual(len(expected_files), len(files1))
        self.assertEqual(len(expected_files), len(files2))

        self.assertEqual(
            expected_files['site_model_file'].id, files1['site_model_file'].id)
        self.assertEqual(
            expected_files['site_model_file'].id, files2['site_model_file'].id)

class FileDigestTestCase(unittest.TestCase):
    """Test the _file_digest() function."""

    PATH = helpers.get_data_path("src_model1.dat")

    def test__file_digest(self):
        # Make sure the digest returned by the function matches the one
        # obtained via /usr/bin/md5sum
        if sys.platform == 'darwin':
            expected = subprocess.check_output(["md5", self.PATH]).split()[-1]
        else:
            expected = subprocess.check_output(
                ["md5sum", self.PATH]).split()[0]
        actual = engine2._file_digest(self.PATH)
        self.assertEqual(expected, actual)

class GetContentTypeTestCase(unittest.TestCase):

    def test__get_content_type(self):
        # no file extension
        self.assertEqual('unknown', engine2._get_content_type('/foo/bar/baz'))
        # xml file extension
        self.assertEqual('xml', engine2._get_content_type('/foo/bar/baz.xml'))
        # hdf5 file extension
        self.assertEqual(
            'HDF5', engine2._get_content_type('/foo/bar/baz.HDF5')
        )

class CreateHazardCalculationTestCase(unittest.TestCase):

    def test_create_hazard_calculation(self):
        # Just the bare minimum set of params to satisfy not null constraints
        # in the db.
        params = {
            'base_path': 'path/to/job.ini',
            'calculation_mode': 'classical',
            'region': '1 1 2 2 3 3',
            'width_of_mfd_bin': '1',
            'rupture_mesh_spacing': '1',
            'area_source_discretization': '2',
            'investigation_time': 50,
            'truncation_level': 0,
            'maximum_distance': 200,
            'number_of_logic_tree_samples': 1,
            'intensity_measure_types_and_levels': dict(PGA=[1, 2, 3, 4]),
            'random_seed': 37,
        }

        owner = helpers.default_user()

        site_model = models.Input(digest='123', path='/foo/bar', size=0,
                                  input_type='site_model', owner=owner)
        site_model.save()
        files = [site_model]

        hc = engine2.create_hazard_calculation(owner, params, files)
        # Normalize/clean fields by fetching a fresh copy from the db.
        hc = models.HazardCalculation.objects.get(id=hc.id)

        self.assertEqual(hc.calculation_mode, 'classical')
        self.assertEqual(hc.width_of_mfd_bin, 1.0)
        self.assertEqual(hc.rupture_mesh_spacing, 1.0)
        self.assertEqual(hc.area_source_discretization, 2.0)
        self.assertEqual(hc.investigation_time, 50.0)
        self.assertEqual(hc.truncation_level, 0.0)
        self.assertEqual(hc.maximum_distance, 200.0)

        # Test the input2haz_calc link:
        [inp2hcs] = models.Input2hcalc.objects.filter(
            hazard_calculation=hc.id)

        self.assertEqual(site_model.id, inp2hcs.input.id)

class CreateRiskCalculationTestCase(unittest.TestCase):

    def test_create_risk_calculation(self):
        params = {
            'hazard_output_id': 1,
            'base_path': 'path/to/job.ini',
            'export_dir': '/tmp/xxx',
            'calculation_mode': 'classical',
            # just some sample params
            'lrem_steps_per_interval': 5,
            'conditional_loss_poes': '0.01, 0.02, 0.05',
            'region_constraint': '-0.5 0.5, 0.5 0.5, 0.5 -0.5, -0.5, -0.5',
        }

        owner = helpers.default_user()

        vuln_file = models.Input(digest='123', path='/foo/bar', size=0,
                                 input_type='vulnerability', owner=owner)
        vuln_file.save()
        exposure_file = models.Input(digest='456', path='/foo/baz', size=0,
                                     input_type='exposure', owner=owner)
        exposure_file.save()

        files = [vuln_file, exposure_file]

        rc = engine2.create_risk_calculation(owner, params, files)
        # Normalize/clean fields by fetching a fresh copy from the db.
        rc = models.RiskCalculation.objects.get(id=rc.id)

        self.assertEqual(rc.calculation_mode, 'classical')
        self.assertEqual(rc.lrem_steps_per_interval, 5)
        self.assertEqual(rc.conditional_loss_poes, [0.01, 0.02, 0.05])
        self.assertEqual(
            rc.region_constraint.wkt,
            ('POLYGON ((-0.5000000000000000 0.5000000000000000, '
             '0.5000000000000000 0.5000000000000000, '
             '0.5000000000000000 -0.5000000000000000, '
             '-0.5000000000000000 -0.5000000000000000, '
             '-0.5000000000000000 0.5000000000000000))'))

class ReadJobProfileFromConfigFileTestCase(unittest.TestCase):
    """Integration test for basic engine functions.

    Test reading/generating a hazard job profile from a config file, then run
    through the validation form.
    """

    def test_read_and_validate_hazard_config(self):
        cfg = helpers.demo_file('simple_fault_demo_hazard/job.ini')
        job = engine2.prepare_job(getpass.getuser())
        params, files = engine2.parse_config(open(cfg, 'r'))
        calculation = engine2.create_hazard_calculation(
            job.owner, params, files.values())

        form = validation.ClassicalHazardForm(
            instance=calculation, files=files
        )
        self.assertTrue(form.is_valid())

class OpenquakeCliTestCase(unittest.TestCase):
    """
    Run "openquake --version" as a separate
    process using `subprocess`.
    """

    def test_run_version(self):
        args = [helpers.RUNNER, "--version"]

        print 'Running:', ' '.join(args)  # this is useful for debugging
        return subprocess.check_call(args)

class DeleteHazCalcTestCase(unittest.TestCase):

    def setUp(self):
        hazard_cfg = helpers.demo_file('simple_fault_demo_hazard/job.ini')
        risk_cfg = helpers.demo_file('classical_psha_based_risk/job.ini')

        self.risk_job, _ = helpers.get_risk_job(
            risk_cfg, hazard_cfg,
            output_type='curves', username=getpass.getuser()
        )
        self.hazard_job = self.risk_job.risk_calculation.hazard_output.oq_job

    def test_del_haz_calc(self):
        pass
        import nose; nose.tools.set_trace()

    def test_del_haz_calc_referenced_by_risk_calc(self):
        pass

    def test_del_haz_calc_output_referenced_by_risk_calc(self):
        pass
<|endoftext|>


<commit_before>from __future__ import unicode_literals

import boto.ec2.elb
from boto.ec2.elb.attributes import (
    LbAttributes,
    ConnectionSettingAttribute,
    ConnectionDrainingAttribute,
    AccessLogAttribute,
    CrossZoneLoadBalancingAttribute,
)
from boto.ec2.elb.policies import Policies
from moto.core import BaseBackend

class FakeHealthCheck(object):
    def __init__(self, timeout, healthy_threshold, unhealthy_threshold,
                 interval, target):
        self.timeout = timeout
        self.healthy_threshold = healthy_threshold
        self.unhealthy_threshold = unhealthy_threshold
        self.interval = interval
        self.target = target

class FakeListener(object):
    def __init__(self, load_balancer_port, instance_port, protocol, ssl_certificate_id):
        self.load_balancer_port = load_balancer_port
        self.instance_port = instance_port
        self.protocol = protocol.upper()
        self.ssl_certificate_id = ssl_certificate_id
        self.policy_names = []

    def __repr__(self):
        return "FakeListener(lbp: %s, inp: %s, pro: %s, cid: %s, policies: %s)" % (self.load_balancer_port, self.instance_port, self.protocol, self.ssl_certificate_id, self.policy_names)

class FakeBackend(object):
    def __init__(self, instance_port):
        self.instance_port = instance_port
        self.policy_names = []

    def __repr__(self):
        return "FakeBackend(inp: %s, policies: %s)" % (self.instance_port, self.policy_names)

class FakeLoadBalancer(object):
    def __init__(self, name, zones, ports, scheme='internet-facing',):
        self.name = name
        self.health_check = None
        self.instance_ids = []
        self.zones = zones
        self.listeners = []
        self.backends = []
        self.scheme = scheme
        self.attributes = FakeLoadBalancer.get_default_attributes()
        self.policies = Policies()
        self.policies.other_policies = []
        self.policies.app_cookie_stickiness_policies = []
        self.policies.lb_cookie_stickiness_policies = []

        for port in ports:
            listener = FakeListener(
                protocol=port['protocol'],
                load_balancer_port=port['load_balancer_port'],
                instance_port=port['instance_port'],
                ssl_certificate_id=port.get('sslcertificate_id'),
            )
            self.listeners.append(listener)

            # it is unclear per the AWS documentation as to when or how backend
            # information gets set, so let's guess and set it here *shrug*
            backend = FakeBackend(
                instance_port=port['instance_port'],
            )
            self.backends.append(backend)

    @classmethod
    def create_from_cloudformation_json(cls, resource_name, cloudformation_json, region_name):
        properties = cloudformation_json['Properties']

        elb_backend = elb_backends[region_name]
        new_elb = elb_backend.create_load_balancer(
            name=properties.get('LoadBalancerName', resource_name),
            zones=properties.get('AvailabilityZones'),
            ports=[],
        )

        instance_ids = cloudformation_json.get('Instances', [])
        for instance_id in instance_ids:
            elb_backend.register_instances(new_elb.name, [instance_id])
        return new_elb

    @property
    def physical_resource_id(self):
        return self.name

    def get_cfn_attribute(self, attribute_name):
        from moto.cloudformation.exceptions import UnformattedGetAttTemplateException
        if attribute_name == 'CanonicalHostedZoneName':
            raise NotImplementedError('"Fn::GetAtt" : [ "{0}" , "CanonicalHostedZoneName" ]"')
        elif attribute_name == 'CanonicalHostedZoneNameID':
            raise NotImplementedError('"Fn::GetAtt" : [ "{0}" , "CanonicalHostedZoneNameID" ]"')
        elif attribute_name == 'DNSName':
            raise NotImplementedError('"Fn::GetAtt" : [ "{0}" , "DNSName" ]"')
        elif attribute_name == 'SourceSecurityGroup.GroupName':
            raise NotImplementedError('"Fn::GetAtt" : [ "{0}" , "SourceSecurityGroup.GroupName" ]"')
        elif attribute_name == 'SourceSecurityGroup.OwnerAlias':
            raise NotImplementedError('"Fn::GetAtt" : [ "{0}" , "SourceSecurityGroup.OwnerAlias" ]"')
        raise UnformattedGetAttTemplateException()

    @classmethod
    def get_default_attributes(cls):
        attributes = LbAttributes()

        cross_zone_load_balancing = CrossZoneLoadBalancingAttribute()
        cross_zone_load_balancing.enabled = False
        attributes.cross_zone_load_balancing = cross_zone_load_balancing

        connection_draining = ConnectionDrainingAttribute()
        connection_draining.enabled = False
        attributes.connection_draining = connection_draining

        access_log = AccessLogAttribute()
        access_log.enabled = False
        attributes.access_log = access_log

        connection_settings = ConnectionSettingAttribute()
        connection_settings.idle_timeout = 60
        attributes.connecting_settings = connection_settings

        return attributes

class ELBBackend(BaseBackend):

    def __init__(self):
        self.load_balancers = {}

    def create_load_balancer(self, name, zones, ports, scheme='internet-facing'):
        print scheme

        new_load_balancer = FakeLoadBalancer(name=name, zones=zones, ports=ports, scheme=scheme)
        self.load_balancers[name] = new_load_balancer
        return new_load_balancer

    def create_load_balancer_listeners(self, name, ports):
        balancer = self.load_balancers.get(name, None)
        if balancer:
            for port in ports:
                protocol = port['protocol']
                instance_port = port['instance_port']
                lb_port = port['load_balancer_port']
                ssl_certificate_id = port.get('sslcertificate_id')
                for listener in balancer.listeners:
                    if lb_port == listener.load_balancer_port:
                        break
                else:
                    balancer.listeners.append(FakeListener(lb_port, instance_port, protocol, ssl_certificate_id))

        return balancer

    def describe_load_balancers(self, names):
        balancers = self.load_balancers.values()
        if names:
            return [balancer for balancer in balancers if balancer.name in names]
        else:
            return balancers

    def delete_load_balancer_listeners(self, name, ports):
        balancer = self.load_balancers.get(name, None)
        listeners = []
        if balancer:
            for lb_port in ports:
                for listener in balancer.listeners:
                    if int(lb_port) == int(listener.load_balancer_port):
                        continue
                    else:
                        listeners.append(listener)
        balancer.listeners = listeners
        return balancer

    def delete_load_balancer(self, load_balancer_name):
        self.load_balancers.pop(load_balancer_name, None)

    def get_load_balancer(self, load_balancer_name):
        return self.load_balancers.get(load_balancer_name)

    def configure_health_check(self, load_balancer_name, timeout,
                               healthy_threshold, unhealthy_threshold, interval,
                               target):
        check = FakeHealthCheck(timeout, healthy_threshold, unhealthy_threshold,
                                interval, target)
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.health_check = check
        return check

    def set_load_balancer_listener_sslcertificate(self, name, lb_port, ssl_certificate_id):
        balancer = self.load_balancers.get(name, None)
        if balancer:
            for idx, listener in enumerate(balancer.listeners):
                if lb_port == listener.load_balancer_port:
                    balancer.listeners[idx].ssl_certificate_id = ssl_certificate_id

        return balancer

    def register_instances(self, load_balancer_name, instance_ids):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.instance_ids.extend(instance_ids)
        return load_balancer

    def deregister_instances(self, load_balancer_name, instance_ids):
        load_balancer = self.get_load_balancer(load_balancer_name)
        new_instance_ids = [instance_id for instance_id in load_balancer.instance_ids if instance_id not in instance_ids]
        load_balancer.instance_ids = new_instance_ids
        return load_balancer

    def set_cross_zone_load_balancing_attribute(self, load_balancer_name, attribute):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.attributes.cross_zone_load_balancing = attribute
        return load_balancer

    def set_access_log_attribute(self, load_balancer_name, attribute):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.attributes.access_log = attribute
        return load_balancer

    def set_connection_draining_attribute(self, load_balancer_name, attribute):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.attributes.connection_draining = attribute
        return load_balancer

    def set_connection_settings_attribute(self, load_balancer_name, attribute):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.attributes.connecting_settings = attribute
        return load_balancer

    def create_lb_other_policy(self, load_balancer_name, other_policy):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.policies.other_policies.append(other_policy)
        return load_balancer

    def create_app_cookie_stickiness_policy(self, load_balancer_name, policy):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.policies.app_cookie_stickiness_policies.append(policy)
        return load_balancer

    def create_lb_cookie_stickiness_policy(self, load_balancer_name, policy):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.policies.lb_cookie_stickiness_policies.append(policy)
        return load_balancer

    def set_load_balancer_policies_of_backend_server(self, load_balancer_name, instance_port, policies):
        load_balancer = self.get_load_balancer(load_balancer_name)
        backend = [b for b in load_balancer.backends if int(b.instance_port) == instance_port][0]
        backend_idx = load_balancer.backends.index(backend)
        backend.policy_names = policies
        load_balancer.backends[backend_idx] = backend
        return load_balancer

    def set_load_balancer_policies_of_listener(self, load_balancer_name, load_balancer_port, policies):
        load_balancer = self.get_load_balancer(load_balancer_name)
        listener = [l for l in load_balancer.listeners if int(l.load_balancer_port) == load_balancer_port][0]
        listener_idx = load_balancer.listeners.index(listener)
        listener.policy_names = policies
        load_balancer.listeners[listener_idx] = listener
        return load_balancer

elb_backends = {}
for region in boto.ec2.elb.regions():
    elb_backends[region.name] = ELBBackend()
<commit_msg>Remove debug print statements<commit_after>from __future__ import unicode_literals

import boto.ec2.elb
from boto.ec2.elb.attributes import (
    LbAttributes,
    ConnectionSettingAttribute,
    ConnectionDrainingAttribute,
    AccessLogAttribute,
    CrossZoneLoadBalancingAttribute,
)
from boto.ec2.elb.policies import Policies
from moto.core import BaseBackend

class FakeHealthCheck(object):
    def __init__(self, timeout, healthy_threshold, unhealthy_threshold,
                 interval, target):
        self.timeout = timeout
        self.healthy_threshold = healthy_threshold
        self.unhealthy_threshold = unhealthy_threshold
        self.interval = interval
        self.target = target

class FakeListener(object):
    def __init__(self, load_balancer_port, instance_port, protocol, ssl_certificate_id):
        self.load_balancer_port = load_balancer_port
        self.instance_port = instance_port
        self.protocol = protocol.upper()
        self.ssl_certificate_id = ssl_certificate_id
        self.policy_names = []

    def __repr__(self):
        return "FakeListener(lbp: %s, inp: %s, pro: %s, cid: %s, policies: %s)" % (self.load_balancer_port, self.instance_port, self.protocol, self.ssl_certificate_id, self.policy_names)

class FakeBackend(object):
    def __init__(self, instance_port):
        self.instance_port = instance_port
        self.policy_names = []

    def __repr__(self):
        return "FakeBackend(inp: %s, policies: %s)" % (self.instance_port, self.policy_names)

class FakeLoadBalancer(object):
    def __init__(self, name, zones, ports, scheme='internet-facing',):
        self.name = name
        self.health_check = None
        self.instance_ids = []
        self.zones = zones
        self.listeners = []
        self.backends = []
        self.scheme = scheme
        self.attributes = FakeLoadBalancer.get_default_attributes()
        self.policies = Policies()
        self.policies.other_policies = []
        self.policies.app_cookie_stickiness_policies = []
        self.policies.lb_cookie_stickiness_policies = []

        for port in ports:
            listener = FakeListener(
                protocol=port['protocol'],
                load_balancer_port=port['load_balancer_port'],
                instance_port=port['instance_port'],
                ssl_certificate_id=port.get('sslcertificate_id'),
            )
            self.listeners.append(listener)

            # it is unclear per the AWS documentation as to when or how backend
            # information gets set, so let's guess and set it here *shrug*
            backend = FakeBackend(
                instance_port=port['instance_port'],
            )
            self.backends.append(backend)

    @classmethod
    def create_from_cloudformation_json(cls, resource_name, cloudformation_json, region_name):
        properties = cloudformation_json['Properties']

        elb_backend = elb_backends[region_name]
        new_elb = elb_backend.create_load_balancer(
            name=properties.get('LoadBalancerName', resource_name),
            zones=properties.get('AvailabilityZones'),
            ports=[],
        )

        instance_ids = cloudformation_json.get('Instances', [])
        for instance_id in instance_ids:
            elb_backend.register_instances(new_elb.name, [instance_id])
        return new_elb

    @property
    def physical_resource_id(self):
        return self.name

    def get_cfn_attribute(self, attribute_name):
        from moto.cloudformation.exceptions import UnformattedGetAttTemplateException
        if attribute_name == 'CanonicalHostedZoneName':
            raise NotImplementedError('"Fn::GetAtt" : [ "{0}" , "CanonicalHostedZoneName" ]"')
        elif attribute_name == 'CanonicalHostedZoneNameID':
            raise NotImplementedError('"Fn::GetAtt" : [ "{0}" , "CanonicalHostedZoneNameID" ]"')
        elif attribute_name == 'DNSName':
            raise NotImplementedError('"Fn::GetAtt" : [ "{0}" , "DNSName" ]"')
        elif attribute_name == 'SourceSecurityGroup.GroupName':
            raise NotImplementedError('"Fn::GetAtt" : [ "{0}" , "SourceSecurityGroup.GroupName" ]"')
        elif attribute_name == 'SourceSecurityGroup.OwnerAlias':
            raise NotImplementedError('"Fn::GetAtt" : [ "{0}" , "SourceSecurityGroup.OwnerAlias" ]"')
        raise UnformattedGetAttTemplateException()

    @classmethod
    def get_default_attributes(cls):
        attributes = LbAttributes()

        cross_zone_load_balancing = CrossZoneLoadBalancingAttribute()
        cross_zone_load_balancing.enabled = False
        attributes.cross_zone_load_balancing = cross_zone_load_balancing

        connection_draining = ConnectionDrainingAttribute()
        connection_draining.enabled = False
        attributes.connection_draining = connection_draining

        access_log = AccessLogAttribute()
        access_log.enabled = False
        attributes.access_log = access_log

        connection_settings = ConnectionSettingAttribute()
        connection_settings.idle_timeout = 60
        attributes.connecting_settings = connection_settings

        return attributes

class ELBBackend(BaseBackend):

    def __init__(self):
        self.load_balancers = {}

    def create_load_balancer(self, name, zones, ports, scheme='internet-facing'):
        new_load_balancer = FakeLoadBalancer(name=name, zones=zones, ports=ports, scheme=scheme)
        self.load_balancers[name] = new_load_balancer
        return new_load_balancer

    def create_load_balancer_listeners(self, name, ports):
        balancer = self.load_balancers.get(name, None)
        if balancer:
            for port in ports:
                protocol = port['protocol']
                instance_port = port['instance_port']
                lb_port = port['load_balancer_port']
                ssl_certificate_id = port.get('sslcertificate_id')
                for listener in balancer.listeners:
                    if lb_port == listener.load_balancer_port:
                        break
                else:
                    balancer.listeners.append(FakeListener(lb_port, instance_port, protocol, ssl_certificate_id))

        return balancer

    def describe_load_balancers(self, names):
        balancers = self.load_balancers.values()
        if names:
            return [balancer for balancer in balancers if balancer.name in names]
        else:
            return balancers

    def delete_load_balancer_listeners(self, name, ports):
        balancer = self.load_balancers.get(name, None)
        listeners = []
        if balancer:
            for lb_port in ports:
                for listener in balancer.listeners:
                    if int(lb_port) == int(listener.load_balancer_port):
                        continue
                    else:
                        listeners.append(listener)
        balancer.listeners = listeners
        return balancer

    def delete_load_balancer(self, load_balancer_name):
        self.load_balancers.pop(load_balancer_name, None)

    def get_load_balancer(self, load_balancer_name):
        return self.load_balancers.get(load_balancer_name)

    def configure_health_check(self, load_balancer_name, timeout,
                               healthy_threshold, unhealthy_threshold, interval,
                               target):
        check = FakeHealthCheck(timeout, healthy_threshold, unhealthy_threshold,
                                interval, target)
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.health_check = check
        return check

    def set_load_balancer_listener_sslcertificate(self, name, lb_port, ssl_certificate_id):
        balancer = self.load_balancers.get(name, None)
        if balancer:
            for idx, listener in enumerate(balancer.listeners):
                if lb_port == listener.load_balancer_port:
                    balancer.listeners[idx].ssl_certificate_id = ssl_certificate_id

        return balancer

    def register_instances(self, load_balancer_name, instance_ids):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.instance_ids.extend(instance_ids)
        return load_balancer

    def deregister_instances(self, load_balancer_name, instance_ids):
        load_balancer = self.get_load_balancer(load_balancer_name)
        new_instance_ids = [instance_id for instance_id in load_balancer.instance_ids if instance_id not in instance_ids]
        load_balancer.instance_ids = new_instance_ids
        return load_balancer

    def set_cross_zone_load_balancing_attribute(self, load_balancer_name, attribute):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.attributes.cross_zone_load_balancing = attribute
        return load_balancer

    def set_access_log_attribute(self, load_balancer_name, attribute):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.attributes.access_log = attribute
        return load_balancer

    def set_connection_draining_attribute(self, load_balancer_name, attribute):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.attributes.connection_draining = attribute
        return load_balancer

    def set_connection_settings_attribute(self, load_balancer_name, attribute):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.attributes.connecting_settings = attribute
        return load_balancer

    def create_lb_other_policy(self, load_balancer_name, other_policy):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.policies.other_policies.append(other_policy)
        return load_balancer

    def create_app_cookie_stickiness_policy(self, load_balancer_name, policy):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.policies.app_cookie_stickiness_policies.append(policy)
        return load_balancer

    def create_lb_cookie_stickiness_policy(self, load_balancer_name, policy):
        load_balancer = self.get_load_balancer(load_balancer_name)
        load_balancer.policies.lb_cookie_stickiness_policies.append(policy)
        return load_balancer

    def set_load_balancer_policies_of_backend_server(self, load_balancer_name, instance_port, policies):
        load_balancer = self.get_load_balancer(load_balancer_name)
        backend = [b for b in load_balancer.backends if int(b.instance_port) == instance_port][0]
        backend_idx = load_balancer.backends.index(backend)
        backend.policy_names = policies
        load_balancer.backends[backend_idx] = backend
        return load_balancer

    def set_load_balancer_policies_of_listener(self, load_balancer_name, load_balancer_port, policies):
        load_balancer = self.get_load_balancer(load_balancer_name)
        listener = [l for l in load_balancer.listeners if int(l.load_balancer_port) == load_balancer_port][0]
        listener_idx = load_balancer.listeners.index(listener)
        listener.policy_names = policies
        load_balancer.listeners[listener_idx] = listener
        return load_balancer

elb_backends = {}
for region in boto.ec2.elb.regions():
    elb_backends[region.name] = ELBBackend()
<|endoftext|>