<commit_before>from django.test import TestCase

from translator.forms import TranslationForm
from translator.models import FailedTranslation
from translator.utils import Translator

class TranslationFormTests(TestCase):
    def test_clean_command(self):
        f = TranslationForm({"command": ""})
        self.assertFalse(f.is_valid())
        self.assertEqual(f.errors["command"], ["This field is required."])

        f = TranslationForm({"command": "bzr commit", "vcs": "git"})
        self.assertFalse(f.is_valid())
        self.assertEqual(f.errors["command"], ["Command must start with a valid VCS (git, hg, svn)."])

        f = TranslationForm({"command": "svn commit", "vcs": "git"})
        self.assertTrue(f.is_valid())

class TranslatorTests(TestCase):
    def assert_translates(self, translator, command, result):
        r = translator.translate(command)
        self.assertTrue(r.success)
        self.assertEqual(r.result, result)

    def assert_cant_handle(self, translator, command):
        r = translator.translate(command)
        self.assertFalse(r.success)
        self.assertTrue(r.result.startswith("We can't handle this yet"))

    def test_x_to_x(self):
        t = Translator("svn", "svn")
        self.assert_translates(t, "log", "svn log")

    def test_svn_to_git(self):
        t = Translator("svn", "git")
        self.assert_translates(t, "commit", "git commit -a && git push")
        self.assert_translates(t, "ci", "git commit -a && git push")
        self.assert_translates(t, "checkout", "git clone")
        self.assert_translates(t, "co", "git clone")
        self.assert_translates(t, "add", "git add")
        self.assert_translates(t, "add file.txt", "git add file.txt")
        self.assert_translates(t, "add some/other/file.txt", "git add some/other/file.txt")
        self.assert_translates(t, "update", "git pull")

    def test_git_to_svn(self):
         t = Translator("git", "svn")
         self.assert_translates(t, "pull", "svn up")
         self.assert_translates(t, "clone", "svn checkout")
         self.assert_translates(t, "status", "svn status")

    def test_hg_to_git(self):
        t = Translator("hg", "git")
        self.assert_translates(t, "pull", "git fetch")
        self.assert_translates(t, "commit", "git commit -a")
        self.assert_translates(t, "push", "git push")
        self.assert_translates(t, "diff", "git diff")

    def test_git_to_hg(self):
        t = Translator("git", "hg")
        self.assert_translates(t, "init", "hg init")
        self.assert_translates(t, "clone", "hg clone")
        self.assert_translates(t, "status", "hg status")
        self.assert_translates(t, "pull", "hg pull -u")
        self.assert_translates(t, "push", "hg push")
        self.assert_translates(t, "diff", "hg diff")

    def test_svn_to_hg(self):
        t = Translator("svn", "hg")
        self.assert_translates(t, "commit", "hg commit && hg push")
        self.assert_translates(t, "checkout", "hg clone")

    def test_cant_handle(self):
        t = Translator("svn", "git")
        self.assert_cant_handle(t, "commit some/file")
        f = FailedTranslation.objects.get()
        self.assertEqual(f.source, "svn")
        self.assertEqual(f.target, "git")
        self.assertEqual(f.command, "commit some/file")
        self.assertEqual(f.count, 1)

        self.assert_cant_handle(t, "commit some/file")
        f = FailedTranslation.objects.get()
        self.assertEqual(f.count, 2)

        t = Translator("git", "svn")
        self.assert_cant_handle(t, "commit -a")

        t = Translator("svn", "hg")
        self.assert_cant_handle(t, "commit -a")
<commit_msg>Some renaming.<commit_after>from django.test import TestCase

from translator.forms import TranslationForm
from translator.models import FailedTranslation
from translator.utils import Translator

class TranslationFormTests(TestCase):
    def test_clean_command(self):
        f = TranslationForm({"command": ""})
        self.assertFalse(f.is_valid())
        self.assertEqual(f.errors["command"], ["This field is required."])

        f = TranslationForm({"command": "bzr commit", "vcs": "git"})
        self.assertFalse(f.is_valid())
        self.assertEqual(f.errors["command"], ["Command must start with a valid VCS (git, hg, svn)."])

        f = TranslationForm({"command": "svn commit", "vcs": "git"})
        self.assertTrue(f.is_valid())

class TranslatorTests(TestCase):
    def assert_translates(self, translator, command, result):
        r = translator.translate(command)
        self.assertTrue(r.success)
        self.assertEqual(r.result, result)

    def assert_cant_handle_yet(self, translator, command):
        r = translator.translate(command)
        self.assertFalse(r.success)
        self.assertTrue(r.result.startswith("We can't handle this yet"))

    def test_x_to_x(self):
        t = Translator("svn", "svn")
        self.assert_translates(t, "log", "svn log")

    def test_svn_to_git(self):
        t = Translator("svn", "git")
        self.assert_translates(t, "commit", "git commit -a && git push")
        self.assert_translates(t, "ci", "git commit -a && git push")
        self.assert_translates(t, "checkout", "git clone")
        self.assert_translates(t, "co", "git clone")
        self.assert_translates(t, "add", "git add")
        self.assert_translates(t, "add file.txt", "git add file.txt")
        self.assert_translates(t, "add some/other/file.txt", "git add some/other/file.txt")
        self.assert_translates(t, "update", "git pull")

    def test_git_to_svn(self):
         t = Translator("git", "svn")
         self.assert_translates(t, "pull", "svn up")
         self.assert_translates(t, "clone", "svn checkout")
         self.assert_translates(t, "status", "svn status")

    def test_hg_to_git(self):
        t = Translator("hg", "git")
        self.assert_translates(t, "pull", "git fetch")
        self.assert_translates(t, "commit", "git commit -a")
        self.assert_translates(t, "push", "git push")
        self.assert_translates(t, "diff", "git diff")

    def test_git_to_hg(self):
        t = Translator("git", "hg")
        self.assert_translates(t, "init", "hg init")
        self.assert_translates(t, "clone", "hg clone")
        self.assert_translates(t, "status", "hg status")
        self.assert_translates(t, "pull", "hg pull -u")
        self.assert_translates(t, "push", "hg push")
        self.assert_translates(t, "diff", "hg diff")

    def test_svn_to_hg(self):
        t = Translator("svn", "hg")
        self.assert_translates(t, "commit", "hg commit && hg push")
        self.assert_translates(t, "checkout", "hg clone")

    def test_cant_handle_yes(self):
        t = Translator("svn", "git")
        self.assert_cant_handle_yet(t, "commit some/file")
        f = FailedTranslation.objects.get()
        self.assertEqual(f.source, "svn")
        self.assertEqual(f.target, "git")
        self.assertEqual(f.command, "commit some/file")
        self.assertEqual(f.count, 1)

        self.assert_cant_handle_yet(t, "commit some/file")
        f = FailedTranslation.objects.get()
        self.assertEqual(f.count, 2)

        t = Translator("git", "svn")
        self.assert_cant_handle_yet(t, "commit -a")

        t = Translator("svn", "hg")
        self.assert_cant_handle_yet(t, "commit -a")
<|endoftext|>


<commit_before>"""
Test lldb data formatter subsystem.
"""

from __future__ import print_function

import os
import time
import re
import lldb
from lldbsuite.test.decorators import *
from lldbsuite.test.lldbtest import *
from lldbsuite.test import lldbutil

class LibcxxListDataFormatterTestCase(TestBase):

    mydir = TestBase.compute_mydir(__file__)

    def setUp(self):
        # Call super's setUp().
        TestBase.setUp(self)
        # Find the line number to break at.
        self.line = line_number('main.cpp', '// Set break point at this line.')
        self.line2 = line_number('main.cpp',
                                 '// Set second break point at this line.')
        self.line3 = line_number('main.cpp',
                                 '// Set third break point at this line.')
        self.line4 = line_number('main.cpp',
                                 '// Set fourth break point at this line.')

    @add_test_categories(["libc++"])
    @skipIf(debug_info="gmodules",
            bugnumber="https://bugs.llvm.org/show_bug.cgi?id=36048")
    def test_with_run_command(self):
        """Test that that file and class static variables display correctly."""
        self.build()
        self.runCmd("file " + self.getBuildArtifact("a.out"), CURRENT_EXECUTABLE_SET)

        lldbutil.run_break_set_by_file_and_line(
            self, "main.cpp", self.line, num_expected_locations=-1)
        lldbutil.run_break_set_by_file_and_line(
            self, "main.cpp", self.line2, num_expected_locations=-1)
        lldbutil.run_break_set_by_file_and_line(
            self, "main.cpp", self.line3, num_expected_locations=-1)
        lldbutil.run_break_set_by_file_and_line(
            self, "main.cpp", self.line4, num_expected_locations=-1)

        self.runCmd("run", RUN_SUCCEEDED)

        lldbutil.skip_if_library_missing(
            self, self.target(), lldbutil.PrintableRegex("libc\+\+"))

        # The stop reason of the thread should be breakpoint.
        self.expect("thread list", STOPPED_DUE_TO_BREAKPOINT,
                    substrs=['stopped',
                             'stop reason = breakpoint'])

        # This is the function to remove the custom formats in order to have a
        # clean slate for the next test case.
        def cleanup():
            self.runCmd('type format clear', check=False)
            self.runCmd('type summary clear', check=False)
            self.runCmd('type filter clear', check=False)
            self.runCmd('type synth clear', check=False)
            self.runCmd(
                "settings set target.max-children-count 256",
                check=False)

        # Execute the cleanup function during test case tear down.
        self.addTearDownHook(cleanup)

        self.runCmd("frame variable numbers_list --show-types")
        self.runCmd(
            "type summary add std::int_list std::string_list int_list string_list --summary-string \"list has ${svar%#} items\" -e")
        self.runCmd("type format add -f hex int")

        self.expect("frame variable numbers_list --raw", matching=False,
                    substrs=['list has 0 items',
                             '{}'])

        self.expect("frame variable numbers_list",
                    substrs=['list has 0 items',
                             '{}'])

        self.expect("p numbers_list",
                    substrs=['list has 0 items',
                             '{}'])

        self.runCmd("n") # This gets up past the printf
        self.runCmd("n") # Now advance over the first push_back.

        self.expect("frame variable numbers_list",
                    substrs=['list has 1 items',
                             '[0] = ',
                             '0x12345678'])

        self.runCmd("n")
        self.runCmd("n")
        self.runCmd("n")

        self.expect("frame variable numbers_list",
                    substrs=['list has 4 items',
                             '[0] = ',
                             '0x12345678',
                             '[1] =',
                             '0x11223344',
                             '[2] =',
                             '0xbeeffeed',
                             '[3] =',
                             '0x00abba00'])

        self.runCmd("n")
        self.runCmd("n")

        self.expect("frame variable numbers_list",
                    substrs=['list has 6 items',
                             '[0] = ',
                             '0x12345678',
                             '0x11223344',
                             '0xbeeffeed',
                             '0x00abba00',
                             '[4] =',
                             '0x0abcdef0',
                             '[5] =',
                             '0x0cab0cab'])

        self.expect("p numbers_list",
                    substrs=['list has 6 items',
                             '[0] = ',
                             '0x12345678',
                             '0x11223344',
                             '0xbeeffeed',
                             '0x00abba00',
                             '[4] =',
                             '0x0abcdef0',
                             '[5] =',
                             '0x0cab0cab'])

        # check access-by-index
        self.expect("frame variable numbers_list[0]",
                    substrs=['0x12345678'])
        self.expect("frame variable numbers_list[1]",
                    substrs=['0x11223344'])

        self.runCmd("n")

        self.expect("frame variable numbers_list",
                    substrs=['list has 0 items',
                             '{}'])

        self.runCmd("n")
        self.runCmd("n")
        self.runCmd("n")
        self.runCmd("n")

        self.expect("frame variable numbers_list",
                    substrs=['list has 4 items',
                             '[0] = ', '1',
                             '[1] = ', '2',
                             '[2] = ', '3',
                             '[3] = ', '4'])

        # check that MightHaveChildren() gets it right
        self.assertTrue(
            self.frame().FindVariable("numbers_list").MightHaveChildren(),
            "numbers_list.MightHaveChildren() says False for non empty!")

        self.runCmd("type format delete int")

        self.runCmd("c")

        self.expect("frame variable text_list",
                    substrs=['list has 3 items',
                             '[0]', 'goofy',
                             '[1]', 'is',
                             '[2]', 'smart'])

        # check that MightHaveChildren() gets it right
        self.assertTrue(
            self.frame().FindVariable("text_list").MightHaveChildren(),
            "text_list.MightHaveChildren() says False for non empty!")

        self.expect("p text_list",
                    substrs=['list has 3 items',
                             '\"goofy\"',
                             '\"is\"',
                             '\"smart\"'])

        self.runCmd("n") # This gets us past the printf
        self.runCmd("n")

        # check access-by-index
        self.expect("frame variable text_list[0]",
                    substrs=['goofy'])
        self.expect("frame variable text_list[3]",
                    substrs=['!!!'])

        self.runCmd("continue")

        # check that the list provider correctly updates if elements move
        countingList = self.frame().FindVariable("countingList")
        countingList.SetPreferDynamicValue(True)
        countingList.SetPreferSyntheticValue(True)

        self.assertTrue(countingList.GetChildAtIndex(
            0).GetValueAsUnsigned(0) == 3141, "list[0] == 3141")
        self.assertTrue(countingList.GetChildAtIndex(
            1).GetValueAsUnsigned(0) == 3141, "list[1] == 3141")

        self.runCmd("continue")

        self.assertTrue(
            countingList.GetChildAtIndex(0).GetValueAsUnsigned(0) == 3141,
            "uniqued list[0] == 3141")
        self.assertTrue(
            countingList.GetChildAtIndex(1).GetValueAsUnsigned(0) == 3142,
            "uniqued list[1] == 3142")
<commit_msg>Relax a data formatter test<commit_after>"""
Test lldb data formatter subsystem.
"""

from __future__ import print_function

import os
import time
import re
import lldb
from lldbsuite.test.decorators import *
from lldbsuite.test.lldbtest import *
from lldbsuite.test import lldbutil

class LibcxxListDataFormatterTestCase(TestBase):

    mydir = TestBase.compute_mydir(__file__)

    def setUp(self):
        # Call super's setUp().
        TestBase.setUp(self)
        # Find the line number to break at.
        self.line = line_number('main.cpp', '// Set break point at this line.')
        self.line2 = line_number('main.cpp',
                                 '// Set second break point at this line.')
        self.line3 = line_number('main.cpp',
                                 '// Set third break point at this line.')
        self.line4 = line_number('main.cpp',
                                 '// Set fourth break point at this line.')

    @add_test_categories(["libc++"])
    @skipIf(debug_info="gmodules",
            bugnumber="https://bugs.llvm.org/show_bug.cgi?id=36048")
    def test_with_run_command(self):
        """Test that that file and class static variables display correctly."""
        self.build()
        self.runCmd("file " + self.getBuildArtifact("a.out"), CURRENT_EXECUTABLE_SET)

        lldbutil.run_break_set_by_file_and_line(
            self, "main.cpp", self.line, num_expected_locations=-1)
        lldbutil.run_break_set_by_file_and_line(
            self, "main.cpp", self.line2, num_expected_locations=-1)
        lldbutil.run_break_set_by_file_and_line(
            self, "main.cpp", self.line3, num_expected_locations=-1)
        lldbutil.run_break_set_by_file_and_line(
            self, "main.cpp", self.line4, num_expected_locations=-1)

        self.runCmd("run", RUN_SUCCEEDED)

        lldbutil.skip_if_library_missing(
            self, self.target(), lldbutil.PrintableRegex("libc\+\+"))

        # The stop reason of the thread should be breakpoint.
        self.expect("thread list", STOPPED_DUE_TO_BREAKPOINT,
                    substrs=['stopped',
                             'stop reason = breakpoint'])

        # This is the function to remove the custom formats in order to have a
        # clean slate for the next test case.
        def cleanup():
            self.runCmd('type format clear', check=False)
            self.runCmd('type summary clear', check=False)
            self.runCmd('type filter clear', check=False)
            self.runCmd('type synth clear', check=False)
            self.runCmd(
                "settings set target.max-children-count 256",
                check=False)

        # Execute the cleanup function during test case tear down.
        self.addTearDownHook(cleanup)

        self.runCmd("frame variable numbers_list --show-types")
        self.runCmd(
            "type summary add std::int_list std::string_list int_list string_list --summary-string \"list has ${svar%#} items\" -e")
        self.runCmd("type format add -f hex int")

        self.expect("frame variable numbers_list --raw", matching=False,
                    substrs=['list has 0 items',
                             '{}'])

        self.expect("frame variable numbers_list",
                    substrs=['list has 0 items',
                             '{}'])

        self.expect("p numbers_list",
                    substrs=['list has 0 items',
                             '{}'])

        self.runCmd("n") # This gets up past the printf
        self.runCmd("n") # Now advance over the first push_back.

        self.expect("frame variable numbers_list",
                    substrs=['list has 1 items',
                             '[0] = ',
                             '0x12345678'])

        self.runCmd("n")
        self.runCmd("n")
        self.runCmd("n")

        self.expect("frame variable numbers_list",
                    substrs=['list has 4 items',
                             '[0] = ',
                             '0x12345678',
                             '[1] =',
                             '0x11223344',
                             '[2] =',
                             '0xbeeffeed',
                             '[3] =',
                             '0x00abba00'])

        self.runCmd("n")
        self.runCmd("n")

        self.expect("frame variable numbers_list",
                    substrs=['list has 6 items',
                             '[0] = ',
                             '0x12345678',
                             '0x11223344',
                             '0xbeeffeed',
                             '0x00abba00',
                             '[4] =',
                             '0x0abcdef0',
                             '[5] =',
                             '0x0cab0cab'])

        self.expect("p numbers_list",
                    substrs=['list has 6 items',
                             '[0] = ',
                             '0x12345678',
                             '0x11223344',
                             '0xbeeffeed',
                             '0x00abba00',
                             '[4] =',
                             '0x0abcdef0',
                             '[5] =',
                             '0x0cab0cab'])

        # check access-by-index
        self.expect("frame variable numbers_list[0]",
                    substrs=['0x12345678'])
        self.expect("frame variable numbers_list[1]",
                    substrs=['0x11223344'])

        self.runCmd("n")

        self.expect("frame variable numbers_list",
                    substrs=['list has 0 items',
                             '{}'])

        self.runCmd("n")
        self.runCmd("n")
        self.runCmd("n")
        self.runCmd("n")

        self.expect("frame variable numbers_list",
                    substrs=['list has 4 items',
                             '[0] = ', '1',
                             '[1] = ', '2',
                             '[2] = ', '3',
                             '[3] = ', '4'])

        # check that MightHaveChildren() gets it right
        self.assertTrue(
            self.frame().FindVariable("numbers_list").MightHaveChildren(),
            "numbers_list.MightHaveChildren() says False for non empty!")

        self.runCmd("type format delete int")

        self.runCmd("c")

        self.expect("frame variable text_list",
                    substrs=['list has 3 items',
                             '[0]', 'goofy',
                             '[1]', 'is',
                             '[2]', 'smart'])

        # check that MightHaveChildren() gets it right
        self.assertTrue(
            self.frame().FindVariable("text_list").MightHaveChildren(),
            "text_list.MightHaveChildren() says False for non empty!")

        self.expect("p text_list",
                    substrs=['list has 3 items',
                             '\"goofy\"',
                             '\"is\"',
                             '\"smart\"'])

        self.runCmd("n") # This gets us past the printf
        self.runCmd("n")
        self.runCmd("n")

        # check access-by-index
        self.expect("frame variable text_list[0]",
                    substrs=['goofy'])
        self.expect("frame variable text_list[3]",
                    substrs=['!!!'])

        self.runCmd("continue")

        # check that the list provider correctly updates if elements move
        countingList = self.frame().FindVariable("countingList")
        countingList.SetPreferDynamicValue(True)
        countingList.SetPreferSyntheticValue(True)

        self.assertTrue(countingList.GetChildAtIndex(
            0).GetValueAsUnsigned(0) == 3141, "list[0] == 3141")
        self.assertTrue(countingList.GetChildAtIndex(
            1).GetValueAsUnsigned(0) == 3141, "list[1] == 3141")

        self.runCmd("continue")

        self.assertTrue(
            countingList.GetChildAtIndex(0).GetValueAsUnsigned(0) == 3141,
            "uniqued list[0] == 3141")
        self.assertTrue(
            countingList.GetChildAtIndex(1).GetValueAsUnsigned(0) == 3142,
            "uniqued list[1] == 3142")
<|endoftext|>


<commit_before>import os

from flask import Flask, url_for
from flask_login import LoginManager
from flask_sqlalchemy import SQLAlchemy, _SessionSignalEvents
from flexmock import flexmock

import sqlalchemy as sa
from sqlalchemy_continuum import (
    make_versioned, remove_versioning, versioning_manager
)
from sqlalchemy_continuum.plugins import FlaskPlugin
from sqlalchemy_continuum.transaction import TransactionFactory
from tests import (
    TestCase,
    get_driver_name,
    get_dns_from_driver,
    uses_native_versioning
)

class TestFlaskPluginConfiguration(object):
    def test_set_factories(self):
        some_func = lambda: None
        some_other_func = lambda: None
        plugin = FlaskPlugin(
            current_user_id_factory=some_func,
            remote_addr_factory=some_other_func
        )
        assert plugin.current_user_id_factory is some_func
        assert plugin.remote_addr_factory is some_other_func

class TestFlaskPlugin(TestCase):
    plugins = [FlaskPlugin()]
    transaction_cls = TransactionFactory()
    user_cls = 'User'

    def setup_method(self, method):
        TestCase.setup_method(self, method)
        self.app = Flask(__name__)
        self.app.secret_key = 'secret'
        self.app.debug = True
        self.setup_views()
        login_manager = LoginManager()
        login_manager.init_app(self.app)
        self.client = self.app.test_client()
        self.context = self.app.test_request_context()
        self.context.push()

        @login_manager.user_loader
        def load_user(id):
            return self.session.query(self.User).get(id)

    def teardown_method(self, method):
        TestCase.teardown_method(self, method)
        self.context.pop()
        self.context = None
        self.client = None
        self.app = None

    def login(self, user):
        """
        Log in the user returned by :meth:`create_user`.

        :returns: the logged in user
        """
        with self.client.session_transaction() as s:
            s['user_id'] = user.id
        return user

    def logout(self, user=None):
        with self.client.session_transaction() as s:
            s['user_id'] = None

    def create_models(self):
        TestCase.create_models(self)

        class User(self.Model):
            __tablename__ = 'user'
            __versioned__ = {
                'base_classes': (self.Model, )
            }

            id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
            name = sa.Column(sa.Unicode(255), nullable=False)
        self.User = User

    def setup_views(self):
        @self.app.route('/simple-flush')
        def test_simple_flush():
            article = self.Article()
            article.name = u'Some article'
            self.session.add(article)
            self.session.commit()
            return ''

        @self.app.route('/raw-sql-and-flush')
        def test_raw_sql_and_flush():
            self.session.execute(
                "INSERT INTO article (name) VALUES ('some article')"
            )
            article = self.Article()
            article.name = u'Some article'
            self.session.add(article)
            self.session.flush()
            self.session.execute(
                "INSERT INTO article (name) VALUES ('some article')"
            )
            self.session.commit()
            return ''

    def test_versioning_inside_request(self):
        user = self.User(name=u'Rambo')
        self.session.add(user)
        self.session.commit()
        self.login(user)
        self.client.get(url_for('.test_simple_flush'))

        article = self.session.query(self.Article).first()
        tx = article.versions[-1].transaction
        assert tx.user.id == user.id

    def test_raw_sql_and_flush(self):
        user = self.User(name=u'Rambo')
        self.session.add(user)
        self.session.commit()
        self.login(user)
        self.client.get(url_for('.test_raw_sql_and_flush'))
        assert (
            self.session.query(versioning_manager.transaction_cls).count() == 2
        )

class TestFlaskPluginWithoutRequestContext(TestCase):
    plugins = [FlaskPlugin()]
    user_cls = 'User'

    def create_models(self):
        class User(self.Model):
            __tablename__ = 'user'
            __versioned__ = {
                'base_classes': (self.Model, )
            }

            id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
            name = sa.Column(sa.Unicode(255), nullable=False)
        self.User = User

        TestCase.create_models(self)

    def test_versioning_outside_request(self):
        user = self.User(name=u'Rambo')
        self.session.add(user)
        self.session.commit()

class TestFlaskPluginWithFlaskSQLAlchemyExtension(object):
    versioning_strategy = 'validity'

    def create_models(self):
        class User(self.db.Model):
            __tablename__ = 'user'
            __versioned__ = {
                'base_classes': (self.db.Model, )
            }

            id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
            name = sa.Column(sa.Unicode(255), nullable=False)

        class Article(self.db.Model):
            __tablename__ = 'article'
            __versioned__ = {
                'base_classes': (self.db.Model, )
            }

            id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
            name = sa.Column(sa.Unicode(255))

        article_tag = sa.Table(
            'article_tag',
            self.db.Model.metadata,
            sa.Column(
                'article_id',
                sa.Integer,
                sa.ForeignKey('article.id'),
                primary_key=True,
            ),
            sa.Column(
                'tag_id',
                sa.Integer,
                sa.ForeignKey('tag.id'),
                primary_key=True
            )
        )

        class Tag(self.db.Model):
            __tablename__ = 'tag'
            __versioned__ = {
                'base_classes': (self.db.Model, )
            }

            id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
            name = sa.Column(sa.Unicode(255))

        Tag.articles = sa.orm.relationship(
            Article,
            secondary=article_tag,
            backref='tags'
        )

        self.User = User
        self.Article = Article
        self.Tag = Tag

    def setup_method(self, method):
        # Mock the event registering of Flask-SQLAlchemy. Currently there is no
        # way of unregistering Flask-SQLAlchemy event listeners, hence the
        # event listeners would affect other tests.
        flexmock(_SessionSignalEvents).should_receive('register')

        self.db = SQLAlchemy()
        make_versioned()

        versioning_manager.transaction_cls = TransactionFactory()
        versioning_manager.options['native_versioning'] = (
            uses_native_versioning()
        )

        self.create_models()

        sa.orm.configure_mappers()

        self.app = Flask(__name__)
        # self.app.config['SQLALCHEMY_ECHO'] = True
        self.app.config['SQLALCHEMY_DATABASE_URI'] = get_dns_from_driver(
            get_driver_name(os.environ.get('DB', 'sqlite'))
        )
        self.db.init_app(self.app)
        self.app.secret_key = 'secret'
        self.app.debug = True
        self.client = self.app.test_client()
        self.context = self.app.test_request_context()
        self.context.push()
        self.db.create_all()

    def teardown_method(self, method):
        remove_versioning()
        self.db.session.remove()
        self.db.drop_all()
        self.db.session.close_all()
        self.db.engine.dispose()
        self.context.pop()
        self.context = None
        self.client = None
        self.app = None

    def test_version_relations(self):
        article = self.Article()
        article.name = u'Some article'
        article.content = u'Some content'
        self.db.session.add(article)
        self.db.session.commit()
        assert not article.versions[0].tags

    def test_single_insert(self):
        article = self.Article()
        article.name = u'Some article'
        article.content = u'Some content'
        tag = self.Tag(name=u'some tag')
        article.tags.append(tag)
        self.db.session.add(article)
        self.db.session.commit()
        assert len(article.versions[0].tags) == 1

    def test_create_transaction_with_scoped_session(self):
        article = self.Article()
        article.name = u'Some article'
        article.content = u'Some content'
        self.db.session.add(article)
        uow = versioning_manager.unit_of_work(self.db.session)
        transaction = uow.create_transaction(self.db.session)
        assert transaction.id

<commit_msg>Fix the login/logout helpers in the FlaskPlugin tests<commit_after>import os

from flask import Flask, url_for
from flask_login import LoginManager
from flask_sqlalchemy import SQLAlchemy, _SessionSignalEvents
from flexmock import flexmock

import sqlalchemy as sa
from sqlalchemy_continuum import (
    make_versioned, remove_versioning, versioning_manager
)
from sqlalchemy_continuum.plugins import FlaskPlugin
from sqlalchemy_continuum.transaction import TransactionFactory
from tests import (
    TestCase,
    get_driver_name,
    get_dns_from_driver,
    uses_native_versioning
)

class TestFlaskPluginConfiguration(object):
    def test_set_factories(self):
        some_func = lambda: None
        some_other_func = lambda: None
        plugin = FlaskPlugin(
            current_user_id_factory=some_func,
            remote_addr_factory=some_other_func
        )
        assert plugin.current_user_id_factory is some_func
        assert plugin.remote_addr_factory is some_other_func

class TestFlaskPlugin(TestCase):
    plugins = [FlaskPlugin()]
    transaction_cls = TransactionFactory()
    user_cls = 'User'

    def setup_method(self, method):
        TestCase.setup_method(self, method)
        self.app = Flask(__name__)
        self.app.secret_key = 'secret'
        self.app.debug = True
        self.setup_views()
        login_manager = LoginManager()
        login_manager.init_app(self.app)
        self.client = self.app.test_client()
        self.context = self.app.test_request_context()
        self.context.push()

        @login_manager.user_loader
        def load_user(id):
            return self.session.query(self.User).get(id)

    def teardown_method(self, method):
        TestCase.teardown_method(self, method)
        self.context.pop()
        self.context = None
        self.client = None
        self.app = None

    def login(self, user):
        """
        Log in the user returned by :meth:`create_user`.

        :returns: the logged in user
        """
        with self.client.session_transaction() as s:
            s['_user_id'] = user.id
        return user

    def logout(self, user=None):
        with self.client.session_transaction() as s:
            s['_user_id'] = None

    def create_models(self):
        TestCase.create_models(self)

        class User(self.Model):
            __tablename__ = 'user'
            __versioned__ = {
                'base_classes': (self.Model, )
            }

            id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
            name = sa.Column(sa.Unicode(255), nullable=False)
        self.User = User

    def setup_views(self):
        @self.app.route('/simple-flush')
        def test_simple_flush():
            article = self.Article()
            article.name = u'Some article'
            self.session.add(article)
            self.session.commit()
            return ''

        @self.app.route('/raw-sql-and-flush')
        def test_raw_sql_and_flush():
            self.session.execute(
                "INSERT INTO article (name) VALUES ('some article')"
            )
            article = self.Article()
            article.name = u'Some article'
            self.session.add(article)
            self.session.flush()
            self.session.execute(
                "INSERT INTO article (name) VALUES ('some article')"
            )
            self.session.commit()
            return ''

    def test_versioning_inside_request(self):
        user = self.User(name=u'Rambo')
        self.session.add(user)
        self.session.commit()
        self.login(user)
        self.client.get(url_for('.test_simple_flush'))

        article = self.session.query(self.Article).first()
        tx = article.versions[-1].transaction
        assert tx.user.id == user.id

    def test_raw_sql_and_flush(self):
        user = self.User(name=u'Rambo')
        self.session.add(user)
        self.session.commit()
        self.login(user)
        self.client.get(url_for('.test_raw_sql_and_flush'))
        assert (
            self.session.query(versioning_manager.transaction_cls).count() == 2
        )

class TestFlaskPluginWithoutRequestContext(TestCase):
    plugins = [FlaskPlugin()]
    user_cls = 'User'

    def create_models(self):
        class User(self.Model):
            __tablename__ = 'user'
            __versioned__ = {
                'base_classes': (self.Model, )
            }

            id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
            name = sa.Column(sa.Unicode(255), nullable=False)
        self.User = User

        TestCase.create_models(self)

    def test_versioning_outside_request(self):
        user = self.User(name=u'Rambo')
        self.session.add(user)
        self.session.commit()

class TestFlaskPluginWithFlaskSQLAlchemyExtension(object):
    versioning_strategy = 'validity'

    def create_models(self):
        class User(self.db.Model):
            __tablename__ = 'user'
            __versioned__ = {
                'base_classes': (self.db.Model, )
            }

            id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
            name = sa.Column(sa.Unicode(255), nullable=False)

        class Article(self.db.Model):
            __tablename__ = 'article'
            __versioned__ = {
                'base_classes': (self.db.Model, )
            }

            id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
            name = sa.Column(sa.Unicode(255))

        article_tag = sa.Table(
            'article_tag',
            self.db.Model.metadata,
            sa.Column(
                'article_id',
                sa.Integer,
                sa.ForeignKey('article.id'),
                primary_key=True,
            ),
            sa.Column(
                'tag_id',
                sa.Integer,
                sa.ForeignKey('tag.id'),
                primary_key=True
            )
        )

        class Tag(self.db.Model):
            __tablename__ = 'tag'
            __versioned__ = {
                'base_classes': (self.db.Model, )
            }

            id = sa.Column(sa.Integer, autoincrement=True, primary_key=True)
            name = sa.Column(sa.Unicode(255))

        Tag.articles = sa.orm.relationship(
            Article,
            secondary=article_tag,
            backref='tags'
        )

        self.User = User
        self.Article = Article
        self.Tag = Tag

    def setup_method(self, method):
        # Mock the event registering of Flask-SQLAlchemy. Currently there is no
        # way of unregistering Flask-SQLAlchemy event listeners, hence the
        # event listeners would affect other tests.
        flexmock(_SessionSignalEvents).should_receive('register')

        self.db = SQLAlchemy()
        make_versioned()

        versioning_manager.transaction_cls = TransactionFactory()
        versioning_manager.options['native_versioning'] = (
            uses_native_versioning()
        )

        self.create_models()

        sa.orm.configure_mappers()

        self.app = Flask(__name__)
        # self.app.config['SQLALCHEMY_ECHO'] = True
        self.app.config['SQLALCHEMY_DATABASE_URI'] = get_dns_from_driver(
            get_driver_name(os.environ.get('DB', 'sqlite'))
        )
        self.db.init_app(self.app)
        self.app.secret_key = 'secret'
        self.app.debug = True
        self.client = self.app.test_client()
        self.context = self.app.test_request_context()
        self.context.push()
        self.db.create_all()

    def teardown_method(self, method):
        remove_versioning()
        self.db.session.remove()
        self.db.drop_all()
        self.db.session.close_all()
        self.db.engine.dispose()
        self.context.pop()
        self.context = None
        self.client = None
        self.app = None

    def test_version_relations(self):
        article = self.Article()
        article.name = u'Some article'
        article.content = u'Some content'
        self.db.session.add(article)
        self.db.session.commit()
        assert not article.versions[0].tags

    def test_single_insert(self):
        article = self.Article()
        article.name = u'Some article'
        article.content = u'Some content'
        tag = self.Tag(name=u'some tag')
        article.tags.append(tag)
        self.db.session.add(article)
        self.db.session.commit()
        assert len(article.versions[0].tags) == 1

    def test_create_transaction_with_scoped_session(self):
        article = self.Article()
        article.name = u'Some article'
        article.content = u'Some content'
        self.db.session.add(article)
        uow = versioning_manager.unit_of_work(self.db.session)
        transaction = uow.create_transaction(self.db.session)
        assert transaction.id
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import migrations, models

class Migration(migrations.Migration):

    dependencies = [
        ('student', '0009_auto_20170111_0422'),
    ]

    operations = [
        migrations.RunSQL(
            "create unique index email on auth_user (email);",
            "drop index email on auth_user;"
        )
    ]
<commit_msg>No-op a migration that needs to be more clever<commit_after># -*- coding: utf-8 -*-
from __future__ import unicode_literals

from django.db import migrations, models

class Migration(migrations.Migration):

    dependencies = [
        ('student', '0009_auto_20170111_0422'),
    ]

    # This migration was to add a constraint that we lost in the Django
    # 1.4->1.8 upgrade. But since the constraint used to be created, production
    # would already have the constraint even before running the migration, and
    # running the migration would fail. We needed to make the migration
    # idempotent.  Instead of reverting this migration while we did that, we
    # edited it to be a SQL no-op, so that people who had already applied it
    # wouldn't end up with a ghost migration.

    # It had been:
    #
    # migrations.RunSQL(
    #   "create unique index email on auth_user (email);",
    #   "drop index email on auth_user;"
    # )

    operations = [
        migrations.RunSQL(
            # Do nothing:
            "select 1",
            "select 1"
        )
    ]
<|endoftext|>


<commit_before># coding: utf8
from __future__ import unicode_literals
import os
from unittest import TestCase
from tempfile import mktemp
import shutil
import datetime
import time

from httmock import all_requests, response, HTTMock
from mock import patch, Mock

from clldclient.util import b

@all_requests
def clld(url, request):
    res = {
        '/resource/languoid/id/stan1295.json': (
            '<http://glottolog.org/>; rel="canonical"; type="text/html"',
            'application/json',
            b('{"id": "stan1295", "name": "Standard German"}')),
        '/resource/languoid/id/stan1295.rdf': (
            '',
            'application/rdf+xml; charset=utf8',
            b("""\
<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"
         xmlns:dcterms="http://purl.org/dc/terms/">
    <rdf:Description rdf:about="http://glottolog.org/resource/languoid/id/stan1295">
        <dcterms:isReferencedBy rdf:resource="http://glottolog.org/resource/reference/id/7242"/>
    </rdf:Description>
</rdf:RDF>
""")),
    }.get(url.path)
    if res is None:
        return response(404, 'not found', {}, None, 5, request)
    return response(
        200, res[2], {'content-type': res[1], 'link': res[0]}, None, 5, request)

class Tests(TestCase):
    def setUp(self):
        self.tmp = mktemp()

    def tearDown(self):
        if os.path.exists(self.tmp):
            shutil.rmtree(self.tmp, ignore_errors=True)

    def test_Cache(self):
        from clldclient.cache import Cache

        with patch('clldclient.cache.user_cache_dir', Mock(return_value=self.tmp)):
            with HTTMock(clld):
                cache = Cache()
                r1 = cache.get('http://glottolog.org/resource/languoid/id/stan1295.json')
                self.assertEquals(r1.content['id'], 'stan1295')
                self.assertEquals(r1.canonical_url, 'http://glottolog.org/')
                r2 = cache.get('http://glottolog.org/resource/languoid/id/stan1295.json')
                self.assertEquals(r1.created, r2.created)
                cache.drop()
                r2 = cache.get('http://glottolog.org/resource/languoid/id/stan1295.json')
                self.assertNotEquals(r1.created, r2.created)
                self.assertRaises(KeyError, cache.get, 'http://glottolog.org/unknown')
                self.assertEquals(cache.get('http://glottolog.org/unknown', default=1), 1)
                res = cache.get('http://glottolog.org/resource/languoid/id/stan1295.rdf')
                self.assertEquals(res.mimetype, 'application/rdf+xml')
                self.assertEquals(
                    res.canonical_url,
                    b('http://glottolog.org/resource/languoid/id/stan1295.rdf'))
                assert hasattr(res.content, 'triples')
                self.assertEquals(res.links, [])
                cached = {r[0]: r[1] for r in cache.stats()}[b('glottolog.org')]
                self.assertEquals(cached, cache.purge(host=b('glottolog.org')))
                now = datetime.datetime.utcnow()
                time.sleep(0.2)
                cache.get('http://glottolog.org/resource/languoid/id/stan1295.json')
                self.assertEquals(0, cache.purge(before=now, host=b('glottolog.org')))
                self.assertEquals(1, cache.purge(after=now, host=b('glottolog.org')))
<commit_msg>Remove calls to encoding function b()<commit_after># coding: utf8
from __future__ import unicode_literals
import os
from unittest import TestCase
from tempfile import mktemp
import shutil
import datetime
import time

from httmock import all_requests, response, HTTMock
from mock import patch, Mock

from clldclient.util import b

@all_requests
def clld(url, request):
    res = {
        '/resource/languoid/id/stan1295.json': (
            '<http://glottolog.org/>; rel="canonical"; type="text/html"',
            'application/json',
            ('{"id": "stan1295", "name": "Standard German"}')),
        '/resource/languoid/id/stan1295.rdf': (
            '',
            'application/rdf+xml; charset=utf8',
            ("""\
<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"
         xmlns:dcterms="http://purl.org/dc/terms/">
    <rdf:Description rdf:about="http://glottolog.org/resource/languoid/id/stan1295">
        <dcterms:isReferencedBy rdf:resource="http://glottolog.org/resource/reference/id/7242"/>
    </rdf:Description>
</rdf:RDF>
""")),
    }.get(url.path)
    if res is None:
        return response(404, 'not found', {}, None, 5, request)
    return response(
        200, res[2], {'content-type': res[1], 'link': res[0]}, None, 5, request)

class Tests(TestCase):
    def setUp(self):
        self.tmp = mktemp()

    def tearDown(self):
        if os.path.exists(self.tmp):
            shutil.rmtree(self.tmp, ignore_errors=True)

    def test_Cache(self):
        from clldclient.cache import Cache

        with patch('clldclient.cache.user_cache_dir', Mock(return_value=self.tmp)):
            with HTTMock(clld):
                cache = Cache()
                r1 = cache.get('http://glottolog.org/resource/languoid/id/stan1295.json')
                self.assertEquals(r1.content['id'], 'stan1295')
                self.assertEquals(r1.canonical_url, 'http://glottolog.org/')
                r2 = cache.get('http://glottolog.org/resource/languoid/id/stan1295.json')
                self.assertEquals(r1.created, r2.created)
                cache.drop()
                r2 = cache.get('http://glottolog.org/resource/languoid/id/stan1295.json')
                self.assertNotEquals(r1.created, r2.created)
                self.assertRaises(KeyError, cache.get, 'http://glottolog.org/unknown')
                self.assertEquals(cache.get('http://glottolog.org/unknown', default=1), 1)
                res = cache.get('http://glottolog.org/resource/languoid/id/stan1295.rdf')
                self.assertEquals(res.mimetype, 'application/rdf+xml')
                self.assertEquals(
                    res.canonical_url,
                    ('http://glottolog.org/resource/languoid/id/stan1295.rdf'))
                assert hasattr(res.content, 'triples')
                self.assertEquals(res.links, [])
                cached = {r[0]: r[1] for r in cache.stats()}[('glottolog.org')]
                self.assertEquals(cached, cache.purge(host=('glottolog.org')))
                now = datetime.datetime.utcnow()
                time.sleep(0.2)
                cache.get('http://glottolog.org/resource/languoid/id/stan1295.json')
                self.assertEquals(0, cache.purge(before=now, host=('glottolog.org')))
                self.assertEquals(1, cache.purge(after=now, host=('glottolog.org')))
<|endoftext|>


<commit_before>"""
Support for MySensors switches.

For more details about this platform, please refer to the documentation at
https://home-assistant.io/components/switch.mysensors/
"""
import logging
import os

import voluptuous as vol

import homeassistant.helpers.config_validation as cv
from homeassistant.components import mysensors
from homeassistant.components.switch import DOMAIN, SwitchDevice
from homeassistant.config import load_yaml_config_file
from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON

_LOGGER = logging.getLogger(__name__)
DEPENDENCIES = []

ATTR_IR_CODE = 'V_IR_SEND'
SERVICE_SEND_IR_CODE = 'mysensors_send_ir_code'

SEND_IR_CODE_SERVICE_SCHEMA = vol.Schema({
    vol.Optional(ATTR_ENTITY_ID): cv.entity_ids,
    vol.Required(ATTR_IR_CODE): cv.string,
})

def setup_platform(hass, config, add_devices, discovery_info=None):
    """Setup the mysensors platform for switches."""
    # Only act if loaded via mysensors by discovery event.
    # Otherwise gateway is not setup.
    if discovery_info is None:
        return

    gateways = hass.data.get(mysensors.MYSENSORS_GATEWAYS)
    if not gateways:
        return

    for gateway in gateways:
        # Define the S_TYPES and V_TYPES that the platform should handle as
        # states. Map them in a dict of lists.
        pres = gateway.const.Presentation
        set_req = gateway.const.SetReq
        map_sv_types = {
            pres.S_DOOR: [set_req.V_ARMED],
            pres.S_MOTION: [set_req.V_ARMED],
            pres.S_SMOKE: [set_req.V_ARMED],
            pres.S_LIGHT: [set_req.V_LIGHT],
            pres.S_LOCK: [set_req.V_LOCK_STATUS],
            pres.S_IR: [set_req.V_IR_SEND],
        }
        device_class_map = {
            pres.S_DOOR: MySensorsSwitch,
            pres.S_MOTION: MySensorsSwitch,
            pres.S_SMOKE: MySensorsSwitch,
            pres.S_LIGHT: MySensorsSwitch,
            pres.S_LOCK: MySensorsSwitch,
            pres.S_IR: MySensorsIRSwitch,
        }
        if float(gateway.protocol_version) >= 1.5:
            map_sv_types.update({
                pres.S_BINARY: [set_req.V_STATUS, set_req.V_LIGHT],
                pres.S_SPRINKLER: [set_req.V_STATUS],
                pres.S_WATER_LEAK: [set_req.V_ARMED],
                pres.S_SOUND: [set_req.V_ARMED],
                pres.S_VIBRATION: [set_req.V_ARMED],
                pres.S_MOISTURE: [set_req.V_ARMED],
            })
            map_sv_types[pres.S_LIGHT].append(set_req.V_STATUS)
            device_class_map.update({
                pres.S_BINARY: MySensorsSwitch,
                pres.S_SPRINKLER: MySensorsSwitch,
                pres.S_WATER_LEAK: MySensorsSwitch,
                pres.S_SOUND: MySensorsSwitch,
                pres.S_VIBRATION: MySensorsSwitch,
                pres.S_MOISTURE: MySensorsSwitch,
            })
        if float(gateway.protocol_version) >= 2.0:
            map_sv_types.update({
                pres.S_WATER_QUALITY: [set_req.V_STATUS],
            })
            device_class_map.update({
                pres.S_WATER_QUALITY: MySensorsSwitch,
            })

        devices = {}
        gateway.platform_callbacks.append(mysensors.pf_callback_factory(
            map_sv_types, devices, add_devices, device_class_map))

    def send_ir_code_service(service):
        """Set IR code as device state attribute."""
        entity_ids = service.data.get(ATTR_ENTITY_ID)
        ir_code = service.data.get(ATTR_IR_CODE)

        if entity_ids:
            _devices = [device for device in devices.values()
                        if isinstance(device, MySensorsIRSwitch) and
                        device.entity_id in entity_ids]
        else:
            _devices = [device for device in devices.values()
                        if isinstance(device, MySensorsIRSwitch)]

        kwargs = {ATTR_IR_CODE: ir_code}
        for device in _devices:
            device.turn_on(**kwargs)

    descriptions = load_yaml_config_file(
        os.path.join(os.path.dirname(__file__), 'services.yaml'))

    hass.services.register(DOMAIN, SERVICE_SEND_IR_CODE,
                           send_ir_code_service,
                           descriptions.get(SERVICE_SEND_IR_CODE),
                           schema=SEND_IR_CODE_SERVICE_SCHEMA)

class MySensorsSwitch(mysensors.MySensorsDeviceEntity, SwitchDevice):
    """Representation of the value of a MySensors Switch child node."""

    @property
    def assumed_state(self):
        """Return True if unable to access real state of entity."""
        return self.gateway.optimistic

    @property
    def is_on(self):
        """Return True if switch is on."""
        if self.value_type in self._values:
            return self._values[self.value_type] == STATE_ON
        return False

    def turn_on(self):
        """Turn the switch on."""
        self.gateway.set_child_value(
            self.node_id, self.child_id, self.value_type, 1)
        if self.gateway.optimistic:
            # optimistically assume that switch has changed state
            self._values[self.value_type] = STATE_ON
            self.schedule_update_ha_state()

    def turn_off(self):
        """Turn the switch off."""
        self.gateway.set_child_value(
            self.node_id, self.child_id, self.value_type, 0)
        if self.gateway.optimistic:
            # optimistically assume that switch has changed state
            self._values[self.value_type] = STATE_OFF
            self.schedule_update_ha_state()

class MySensorsIRSwitch(MySensorsSwitch):
    """IR switch child class to MySensorsSwitch."""

    def __init__(self, *args):
        """Setup instance attributes."""
        MySensorsSwitch.__init__(self, *args)
        self._ir_code = None

    @property
    def is_on(self):
        """Return True if switch is on."""
        set_req = self.gateway.const.SetReq
        if set_req.V_LIGHT in self._values:
            return self._values[set_req.V_LIGHT] == STATE_ON
        return False

    def turn_on(self, **kwargs):
        """Turn the IR switch on."""
        set_req = self.gateway.const.SetReq
        if set_req.V_LIGHT not in self._values:
            _LOGGER.error('missing value_type: %s at node: %s, child: %s',
                          set_req.V_LIGHT.name, self.node_id, self.child_id)
            return
        if ATTR_IR_CODE in kwargs:
            self._ir_code = kwargs[ATTR_IR_CODE]
        self.gateway.set_child_value(
            self.node_id, self.child_id, self.value_type, self._ir_code)
        self.gateway.set_child_value(
            self.node_id, self.child_id, set_req.V_LIGHT, 1)
        if self.gateway.optimistic:
            # optimistically assume that switch has changed state
            self._values[self.value_type] = self._ir_code
            self._values[set_req.V_LIGHT] = STATE_ON
            self.schedule_update_ha_state()
            # turn off switch after switch was turned on
            self.turn_off()

    def turn_off(self):
        """Turn the IR switch off."""
        set_req = self.gateway.const.SetReq
        if set_req.V_LIGHT not in self._values:
            _LOGGER.error('missing value_type: %s at node: %s, child: %s',
                          set_req.V_LIGHT.name, self.node_id, self.child_id)
            return
        self.gateway.set_child_value(
            self.node_id, self.child_id, set_req.V_LIGHT, 0)
        if self.gateway.optimistic:
            # optimistically assume that switch has changed state
            self._values[set_req.V_LIGHT] = STATE_OFF
            self.schedule_update_ha_state()

    def update(self):
        """Update the controller with the latest value from a sensor."""
        MySensorsSwitch.update(self)
        if self.value_type in self._values:
            self._ir_code = self._values[self.value_type]
<commit_msg>Fix mysensors ir switch overwriting devices (#4612)<commit_after>"""
Support for MySensors switches.

For more details about this platform, please refer to the documentation at
https://home-assistant.io/components/switch.mysensors/
"""
import logging
import os

import voluptuous as vol

import homeassistant.helpers.config_validation as cv
from homeassistant.components import mysensors
from homeassistant.components.switch import DOMAIN, SwitchDevice
from homeassistant.config import load_yaml_config_file
from homeassistant.const import ATTR_ENTITY_ID, STATE_OFF, STATE_ON

_LOGGER = logging.getLogger(__name__)
DEPENDENCIES = []

ATTR_IR_CODE = 'V_IR_SEND'
SERVICE_SEND_IR_CODE = 'mysensors_send_ir_code'

SEND_IR_CODE_SERVICE_SCHEMA = vol.Schema({
    vol.Optional(ATTR_ENTITY_ID): cv.entity_ids,
    vol.Required(ATTR_IR_CODE): cv.string,
})

def setup_platform(hass, config, add_devices, discovery_info=None):
    """Setup the mysensors platform for switches."""
    # Only act if loaded via mysensors by discovery event.
    # Otherwise gateway is not setup.
    if discovery_info is None:
        return

    gateways = hass.data.get(mysensors.MYSENSORS_GATEWAYS)
    if not gateways:
        return

    platform_devices = []

    for gateway in gateways:
        # Define the S_TYPES and V_TYPES that the platform should handle as
        # states. Map them in a dict of lists.
        pres = gateway.const.Presentation
        set_req = gateway.const.SetReq
        map_sv_types = {
            pres.S_DOOR: [set_req.V_ARMED],
            pres.S_MOTION: [set_req.V_ARMED],
            pres.S_SMOKE: [set_req.V_ARMED],
            pres.S_LIGHT: [set_req.V_LIGHT],
            pres.S_LOCK: [set_req.V_LOCK_STATUS],
            pres.S_IR: [set_req.V_IR_SEND],
        }
        device_class_map = {
            pres.S_DOOR: MySensorsSwitch,
            pres.S_MOTION: MySensorsSwitch,
            pres.S_SMOKE: MySensorsSwitch,
            pres.S_LIGHT: MySensorsSwitch,
            pres.S_LOCK: MySensorsSwitch,
            pres.S_IR: MySensorsIRSwitch,
        }
        if float(gateway.protocol_version) >= 1.5:
            map_sv_types.update({
                pres.S_BINARY: [set_req.V_STATUS, set_req.V_LIGHT],
                pres.S_SPRINKLER: [set_req.V_STATUS],
                pres.S_WATER_LEAK: [set_req.V_ARMED],
                pres.S_SOUND: [set_req.V_ARMED],
                pres.S_VIBRATION: [set_req.V_ARMED],
                pres.S_MOISTURE: [set_req.V_ARMED],
            })
            map_sv_types[pres.S_LIGHT].append(set_req.V_STATUS)
            device_class_map.update({
                pres.S_BINARY: MySensorsSwitch,
                pres.S_SPRINKLER: MySensorsSwitch,
                pres.S_WATER_LEAK: MySensorsSwitch,
                pres.S_SOUND: MySensorsSwitch,
                pres.S_VIBRATION: MySensorsSwitch,
                pres.S_MOISTURE: MySensorsSwitch,
            })
        if float(gateway.protocol_version) >= 2.0:
            map_sv_types.update({
                pres.S_WATER_QUALITY: [set_req.V_STATUS],
            })
            device_class_map.update({
                pres.S_WATER_QUALITY: MySensorsSwitch,
            })

        devices = {}
        gateway.platform_callbacks.append(mysensors.pf_callback_factory(
            map_sv_types, devices, add_devices, device_class_map))
        platform_devices.append(devices)

    def send_ir_code_service(service):
        """Set IR code as device state attribute."""
        entity_ids = service.data.get(ATTR_ENTITY_ID)
        ir_code = service.data.get(ATTR_IR_CODE)

        if entity_ids:
            _devices = [device for gw_devs in platform_devices
                        for device in gw_devs.values()
                        if isinstance(device, MySensorsIRSwitch) and
                        device.entity_id in entity_ids]
        else:
            _devices = [device for gw_devs in platform_devices
                        for device in gw_devs.values()
                        if isinstance(device, MySensorsIRSwitch)]

        kwargs = {ATTR_IR_CODE: ir_code}
        for device in _devices:
            device.turn_on(**kwargs)

    descriptions = load_yaml_config_file(
        os.path.join(os.path.dirname(__file__), 'services.yaml'))

    hass.services.register(DOMAIN, SERVICE_SEND_IR_CODE,
                           send_ir_code_service,
                           descriptions.get(SERVICE_SEND_IR_CODE),
                           schema=SEND_IR_CODE_SERVICE_SCHEMA)

class MySensorsSwitch(mysensors.MySensorsDeviceEntity, SwitchDevice):
    """Representation of the value of a MySensors Switch child node."""

    @property
    def assumed_state(self):
        """Return True if unable to access real state of entity."""
        return self.gateway.optimistic

    @property
    def is_on(self):
        """Return True if switch is on."""
        if self.value_type in self._values:
            return self._values[self.value_type] == STATE_ON
        return False

    def turn_on(self):
        """Turn the switch on."""
        self.gateway.set_child_value(
            self.node_id, self.child_id, self.value_type, 1)
        if self.gateway.optimistic:
            # optimistically assume that switch has changed state
            self._values[self.value_type] = STATE_ON
            self.schedule_update_ha_state()

    def turn_off(self):
        """Turn the switch off."""
        self.gateway.set_child_value(
            self.node_id, self.child_id, self.value_type, 0)
        if self.gateway.optimistic:
            # optimistically assume that switch has changed state
            self._values[self.value_type] = STATE_OFF
            self.schedule_update_ha_state()

class MySensorsIRSwitch(MySensorsSwitch):
    """IR switch child class to MySensorsSwitch."""

    def __init__(self, *args):
        """Setup instance attributes."""
        MySensorsSwitch.__init__(self, *args)
        self._ir_code = None

    @property
    def is_on(self):
        """Return True if switch is on."""
        set_req = self.gateway.const.SetReq
        if set_req.V_LIGHT in self._values:
            return self._values[set_req.V_LIGHT] == STATE_ON
        return False

    def turn_on(self, **kwargs):
        """Turn the IR switch on."""
        set_req = self.gateway.const.SetReq
        if set_req.V_LIGHT not in self._values:
            _LOGGER.error('missing value_type: %s at node: %s, child: %s',
                          set_req.V_LIGHT.name, self.node_id, self.child_id)
            return
        if ATTR_IR_CODE in kwargs:
            self._ir_code = kwargs[ATTR_IR_CODE]
        self.gateway.set_child_value(
            self.node_id, self.child_id, self.value_type, self._ir_code)
        self.gateway.set_child_value(
            self.node_id, self.child_id, set_req.V_LIGHT, 1)
        if self.gateway.optimistic:
            # optimistically assume that switch has changed state
            self._values[self.value_type] = self._ir_code
            self._values[set_req.V_LIGHT] = STATE_ON
            self.schedule_update_ha_state()
            # turn off switch after switch was turned on
            self.turn_off()

    def turn_off(self):
        """Turn the IR switch off."""
        set_req = self.gateway.const.SetReq
        if set_req.V_LIGHT not in self._values:
            _LOGGER.error('missing value_type: %s at node: %s, child: %s',
                          set_req.V_LIGHT.name, self.node_id, self.child_id)
            return
        self.gateway.set_child_value(
            self.node_id, self.child_id, set_req.V_LIGHT, 0)
        if self.gateway.optimistic:
            # optimistically assume that switch has changed state
            self._values[set_req.V_LIGHT] = STATE_OFF
            self.schedule_update_ha_state()

    def update(self):
        """Update the controller with the latest value from a sensor."""
        MySensorsSwitch.update(self)
        if self.value_type in self._values:
            self._ir_code = self._values[self.value_type]
<|endoftext|>


<commit_before>import io
import zipfile
from xml.parsers.expat import ExpatError
from ciso8601 import parse_datetime
from multiprocessing.dummy import Pool
import xmltodict
from django.utils import timezone
from django.contrib.gis.geos import Point
from django.db.models import Q, Exists, OuterRef
from ..import_live_vehicles import ImportLiveVehiclesCommand
from busstops.models import Operator, Service, Locality, StopPoint, ServiceCode
from bustimes.models import Trip
from ...models import Vehicle, VehicleJourney, VehicleLocation

class Command(ImportLiveVehiclesCommand):
    source_name = 'Bus Open Data'
    wait = 20
    cache = set()
    operators = {
        'ASC': ['AKSS', 'ARHE', 'AMTM', 'GLAR'],
        'ANE': ['ARDU', 'ANUM', 'ANEA'],
        'ANW': ['ANWE', 'AMSY', 'ACYM'],
        'ATS': ['ARBB', 'ASES', 'GLAR'],
        'AMD': ['AMID', 'AMNO', 'AFCL'],
        'GOEA': ['KCTB', 'HEDO', 'CHAM'],
        'CBBH': ['CBBH', 'CBNL'],
        'UNO': ['UNOE', 'UNIB'],
        'UNIB': ['UNOE', 'UNIB'],
        'TDY': ['YCST', 'LNUD', 'ROST', 'BPTR', 'KDTR', 'HRGT'],
    }
    operator_cache = {}
    vehicle_cache = {}
    service_cache = {}
    reg_operators = {'BDRB', 'COMT', 'TDY', 'ROST'}
    identifiers = {}

    @staticmethod
    def get_datetime(item):
        return parse_datetime(item['RecordedAtTime'])

    def get_operator(self, operator_ref):
        if operator_ref == "SCEM":
            operator_ref = "SCGH"
        elif operator_ref == "SCSO":
            operator_ref = "SCCO"

        if operator_ref in self.operators:
            return self.operators[operator_ref]

        if operator_ref in self.operator_cache:
            return self.operator_cache[operator_ref]

        condition = Exists(self.source.operatorcode_set.filter(operator=OuterRef('id'), code=operator_ref))
        if len(operator_ref) == 4:
            condition |= Q(id=operator_ref)

        try:
            operator = Operator.objects.get(condition)
            self.operator_cache[operator_ref] = operator
            return operator
        except Operator.DoesNotExist:
            pass

        self.operator_cache[operator_ref] = None

    def get_vehicle(self, item):
        monitored_vehicle_journey = item['MonitoredVehicleJourney']

        operator_ref = monitored_vehicle_journey['OperatorRef']

        vehicle_ref = monitored_vehicle_journey['VehicleRef']
        cache_key = f'{operator_ref}-{vehicle_ref}'
        try:
            return self.vehicles.get(id=self.vehicle_cache[cache_key]), False
        except (KeyError, Vehicle.DoesNotExist):
            pass

        operator = self.get_operator(operator_ref)

        if operator and vehicle_ref.startswith(f'{operator_ref}-'):
            vehicle_ref = vehicle_ref[len(operator_ref) + 1:]

        defaults = {
            'code': vehicle_ref,
            'source': self.source
        }

        if type(operator) is Operator:
            defaults['operator'] = operator
            if operator.parent:
                condition = Q(operator__parent=operator.parent)

                if not vehicle_ref.isdigit() and vehicle_ref.isupper():
                    # Abus operate the 349 using First ticket machines
                    if operator.id == "FBRI":
                        condition |= Q(operator="ABUS")
                    # Connexions Buses 64
                    elif operator.id == "FWYO":
                        condition |= Q(operator="HCTY")

                vehicles = self.vehicles.filter(condition)
            else:
                vehicles = self.vehicles.filter(operator=operator)
        elif type(operator) is list:
            defaults['operator_id'] = operator[0]
            vehicles = self.vehicles.filter(operator__in=operator)
        else:
            vehicles = self.vehicles.filter(operator=None)

        assert vehicle_ref

        condition = Q(code=vehicle_ref)
        if operator:
            if vehicle_ref.isdigit():
                defaults['fleet_number'] = vehicle_ref
                condition |= Q(code__endswith=f'-{vehicle_ref}') | Q(code__startswith=f'{vehicle_ref}_')
            else:
                if '_-_' in vehicle_ref:
                    fleet_number, reg = vehicle_ref.split('_-_', 2)
                    if fleet_number.isdigit():
                        defaults['fleet_number'] = fleet_number
                        reg = reg.replace('_', '')
                        defaults['reg'] = reg
                        if operator_ref in self.reg_operators:
                            condition |= Q(reg=reg)
                elif operator_ref in self.reg_operators:
                    reg = vehicle_ref.replace('_', '')
                    condition |= Q(reg=reg)
                elif operator_ref == 'WHIP':
                    code = vehicle_ref.replace('_', '')
                    condition |= Q(fleet_code=code)
        vehicles = vehicles.filter(condition)

        try:
            vehicle, created = vehicles.get_or_create(defaults)
        except Vehicle.MultipleObjectsReturned as e:
            print(e, operator, vehicle_ref)
            vehicle = vehicles.first()
            created = False

        self.vehicle_cache[cache_key] = vehicle.id
        return vehicle, created

    def get_service(self, operator, item, vehicle_operator_id):
        monitored_vehicle_journey = item['MonitoredVehicleJourney']

        line_ref = monitored_vehicle_journey.get('LineRef')
        if not line_ref:
            return

        destination_ref = monitored_vehicle_journey.get("DestinationRef")

        cache_key = f"{operator}:{vehicle_operator_id}:{line_ref}:{destination_ref}"
        if cache_key in self.service_cache:
            return self.service_cache[cache_key]

        services = Service.objects.filter(
            Exists(ServiceCode.objects.filter(service=OuterRef('id'), scheme__endswith=' SIRI', code=line_ref))
            | Q(line_name__iexact=line_ref),
            current=True
        )

        if type(operator) is Operator and operator.parent and destination_ref:
            services = services.filter(operator__parent=operator.parent)
            # we will use the destination ref to find out exactly which operator it is

        else:
            if type(operator) is Operator:
                condition = Q(operator=operator)
                if vehicle_operator_id != operator.id:
                    condition |= Q(operator=vehicle_operator_id)
                services = services.filter(condition)
            elif type(operator) is list:
                services = services.filter(operator__in=operator)

            try:
                return services.get()
            except Service.DoesNotExist:
                self.service_cache[cache_key] = None
                return
            except Service.MultipleObjectsReturned:
                pass

        if destination_ref:
            try:
                stops = StopPoint.objects.filter(service=OuterRef("pk"), locality__stoppoint=destination_ref)
                return services.filter(Exists(stops)).get()
            except Service.DoesNotExist:
                self.service_cache[cache_key] = None
                return
            except Service.MultipleObjectsReturned:
                try:
                    trips = Trip.objects.filter(route__service=OuterRef("pk"), destination=destination_ref)
                    return services.filter(Exists(trips)).get()
                except (Service.DoesNotExist, Service.MultipleObjectsReturned):
                    pass

        try:
            when = self.get_datetime(item)
            when = when.strftime('%a').lower()
            trips = Trip.objects.filter(**{f'calendar__{when}': True}, route__service=OuterRef("pk"))
            return services.filter(Exists(trips)).get()
        except (Service.DoesNotExist, Service.MultipleObjectsReturned):
            pass

        vehicle_journey_ref = monitored_vehicle_journey.get('VehicleJourneyRef')
        if vehicle_journey_ref and vehicle_journey_ref.isdigit():
            try:
                trips = Trip.objects.filter(route__service=OuterRef("pk"), ticket_machine_code=vehicle_journey_ref)
                return services.filter(Exists(trips)).get()
            except (Service.DoesNotExist, Service.MultipleObjectsReturned):
                pass

    def get_journey(self, item, vehicle):
        monitored_vehicle_journey = item['MonitoredVehicleJourney']

        vehicle_journey_ref = monitored_vehicle_journey.get('VehicleJourneyRef')
        if vehicle_journey_ref == 'UNKNOWN':
            vehicle_journey_ref = None

        route_name = monitored_vehicle_journey.get('LineRef') or ''

        origin_aimed_departure_time = monitored_vehicle_journey.get('OriginAimedDepartureTime')
        if origin_aimed_departure_time:
            origin_aimed_departure_time = parse_datetime(origin_aimed_departure_time)

        journey = None

        # do destination now, in case the route name is in the destination field
        destination = monitored_vehicle_journey.get('DestinationName') or ''
        if vehicle.operator_id == 'TGTC' and destination and not route_name:
            parts = destination.split()
            if parts[0].isdigit() or parts[0][:-1].isdigit():
                route_name = parts[0]
                destination = ' '.join(parts[1:])
                monitored_vehicle_journey['LineRef'] = route_name

        latest_location = vehicle.latest_location
        if latest_location:
            if origin_aimed_departure_time == latest_location.journey.datetime:
                journey = latest_location.journey
            elif vehicle_journey_ref:
                if vehicle_journey_ref != latest_location.journey.code and '_' in vehicle_journey_ref:
                    journey = vehicle.vehiclejourney_set.filter(route_name=route_name, code=vehicle_journey_ref).first()

            if not journey:
                journey = vehicle.vehiclejourney_set.filter(datetime=origin_aimed_departure_time).first()

        if not journey:
            journey = VehicleJourney(
                route_name=route_name,
                vehicle=vehicle,
                source=self.source,
                data=item,
                datetime=origin_aimed_departure_time,
                destination=destination
            )

        if vehicle_journey_ref:
            journey.code = vehicle_journey_ref

        if not journey.destination:
            destination_ref = monitored_vehicle_journey.get('DestinationRef')
            if destination_ref:
                try:
                    journey.destination = Locality.objects.get(stoppoint=destination_ref).name
                except Locality.DoesNotExist:
                    pass

        if (
            latest_location and latest_location.journey.service
            and latest_location.journey.code == journey.code
            and latest_location.journey.route_name == journey.route_name
        ):
            journey.service = latest_location.journey.service
        else:
            operator_ref = monitored_vehicle_journey["OperatorRef"]
            if operator_ref == "FWYO":
                if vehicle.operator_id == "HCTY":
                    operator = ["HCTY"]  # First West Yorkshire/Connexions
                else:
                    operator = ["FWYO", "FLDS"]
            elif operator_ref == "FHAL":
                operator = ["FHAL", "FHUD"]
            elif operator_ref == "FBRI" and vehicle.operator_id == "ABUS":
                operator = ["ABUS"]
            elif operator_ref == "RRAR":
                operator = ["RRAR", "FTVA"]  # Reading RailAir/First Berkshire
            elif operator_ref == "ROST":
                operator = ["ROST", "LNUD", "BPTR"]  # Rosso/Blackburn/Burnley
            else:
                operator = self.get_operator(operator_ref)
            journey.service = self.get_service(operator, item, vehicle.operator_id)

        return journey

    @staticmethod
    def create_vehicle_location(item):
        monitored_vehicle_journey = item['MonitoredVehicleJourney']
        location = monitored_vehicle_journey['VehicleLocation']
        latlong = Point(float(location['Longitude']), float(location['Latitude']))
        return VehicleLocation(
            latlong=latlong,
        )

    def get_changed_items(self, items):
        identifiers = {}
        for item in items:
            monitored_vehicle_journey = item['MonitoredVehicleJourney']
            key = f"{monitored_vehicle_journey['OperatorRef']}-{monitored_vehicle_journey['VehicleRef']}"
            if self.identifiers.get(key) != item['RecordedAtTime']:
                yield item
                identifiers[key] = item['RecordedAtTime']

    def update(self):
        now = timezone.now()

        pool = Pool(8)
        items = self.get_items()

        if items:
            pool.starmap(self.handle_item, ((item, now) for item in self.get_changed_items(items)))
        else:
            return 300  # no items - wait five minutes

        time_taken = timezone.now() - now
        print(time_taken)
        time_taken = time_taken.total_seconds()
        if time_taken < self.wait:
            return self.wait - time_taken
        return 0  # took longer than self.wait

    @staticmethod
    def items_from_response(response):
        return xmltodict.parse(
            response,
            dict_constructor=dict,  # override OrderedDict, cos dict is ordered in modern versions of Python
            force_list=['VehicleActivity']
        )

    def get_items(self):
        response = self.session.get(self.source.url, params=self.source.settings)
        if not response.ok:
            if 'datafeed' in self.source.url:
                print(response.content.decode())
            else:
                print(response)
            return

        if 'datafeed' in self.source.url:
            try:
                data = self.items_from_response(response.content)
            except ExpatError:
                print(response.content.decode())
                return
        else:
            try:
                with zipfile.ZipFile(io.BytesIO(response.content)) as archive:
                    assert archive.namelist() == ['siri.xml']
                    with archive.open('siri.xml') as open_file:
                        try:
                            data = self.items_from_response(open_file)
                        except ExpatError:
                            print(open_file.read())
                            return
            except zipfile.BadZipFile:
                print(response.content.decode())
                return

        self.source.datetime = parse_datetime(data['Siri']['ServiceDelivery']['ResponseTimestamp'])

        return data['Siri']['ServiceDelivery']['VehicleMonitoringDelivery']['VehicleActivity']
<commit_msg>use sirivm worker again<commit_after>import io
import zipfile
import xmltodict
from xml.parsers.expat import ExpatError
from ciso8601 import parse_datetime
from asgiref.sync import async_to_sync
from channels.layers import get_channel_layer
from channels.exceptions import ChannelFull
from django.utils import timezone
from django.contrib.gis.geos import Point
from django.db.models import Q, Exists, OuterRef
from ..import_live_vehicles import ImportLiveVehiclesCommand
from busstops.models import Operator, Service, Locality, StopPoint, ServiceCode
from bustimes.models import Trip
from ...models import Vehicle, VehicleJourney, VehicleLocation

class Command(ImportLiveVehiclesCommand):
    source_name = 'Bus Open Data'
    wait = 20
    cache = set()
    operators = {
        'ASC': ['AKSS', 'ARHE', 'AMTM', 'GLAR'],
        'ANE': ['ARDU', 'ANUM', 'ANEA'],
        'ANW': ['ANWE', 'AMSY', 'ACYM'],
        'ATS': ['ARBB', 'ASES', 'GLAR'],
        'AMD': ['AMID', 'AMNO', 'AFCL'],
        'GOEA': ['KCTB', 'HEDO', 'CHAM'],
        'CBBH': ['CBBH', 'CBNL'],
        'UNO': ['UNOE', 'UNIB'],
        'UNIB': ['UNOE', 'UNIB'],
        'TDY': ['YCST', 'LNUD', 'ROST', 'BPTR', 'KDTR', 'HRGT'],
    }
    operator_cache = {}
    vehicle_cache = {}
    service_cache = {}
    reg_operators = {'BDRB', 'COMT', 'TDY', 'ROST'}
    identifiers = {}

    @staticmethod
    def get_datetime(item):
        return parse_datetime(item['RecordedAtTime'])

    def get_operator(self, operator_ref):
        if operator_ref == "SCEM":
            operator_ref = "SCGH"
        elif operator_ref == "SCSO":
            operator_ref = "SCCO"

        if operator_ref in self.operators:
            return self.operators[operator_ref]

        if operator_ref in self.operator_cache:
            return self.operator_cache[operator_ref]

        condition = Exists(self.source.operatorcode_set.filter(operator=OuterRef('id'), code=operator_ref))
        if len(operator_ref) == 4:
            condition |= Q(id=operator_ref)

        try:
            operator = Operator.objects.get(condition)
            self.operator_cache[operator_ref] = operator
            return operator
        except Operator.DoesNotExist:
            pass

        self.operator_cache[operator_ref] = None

    def get_vehicle(self, item):
        monitored_vehicle_journey = item['MonitoredVehicleJourney']

        operator_ref = monitored_vehicle_journey['OperatorRef']

        vehicle_ref = monitored_vehicle_journey['VehicleRef']
        cache_key = f'{operator_ref}-{vehicle_ref}'
        try:
            return self.vehicles.get(id=self.vehicle_cache[cache_key]), False
        except (KeyError, Vehicle.DoesNotExist):
            pass

        operator = self.get_operator(operator_ref)

        if operator and vehicle_ref.startswith(f'{operator_ref}-'):
            vehicle_ref = vehicle_ref[len(operator_ref) + 1:]

        defaults = {
            'code': vehicle_ref,
            'source': self.source
        }

        if type(operator) is Operator:
            defaults['operator'] = operator
            if operator.parent:
                condition = Q(operator__parent=operator.parent)

                if not vehicle_ref.isdigit() and vehicle_ref.isupper():
                    # Abus operate the 349 using First ticket machines
                    if operator.id == "FBRI":
                        condition |= Q(operator="ABUS")
                    # Connexions Buses 64
                    elif operator.id == "FWYO":
                        condition |= Q(operator="HCTY")

                vehicles = self.vehicles.filter(condition)
            else:
                vehicles = self.vehicles.filter(operator=operator)
        elif type(operator) is list:
            defaults['operator_id'] = operator[0]
            vehicles = self.vehicles.filter(operator__in=operator)
        else:
            vehicles = self.vehicles.filter(operator=None)

        assert vehicle_ref

        condition = Q(code=vehicle_ref)
        if operator:
            if vehicle_ref.isdigit():
                defaults['fleet_number'] = vehicle_ref
                condition |= Q(code__endswith=f'-{vehicle_ref}') | Q(code__startswith=f'{vehicle_ref}_')
            else:
                if '_-_' in vehicle_ref:
                    fleet_number, reg = vehicle_ref.split('_-_', 2)
                    if fleet_number.isdigit():
                        defaults['fleet_number'] = fleet_number
                        reg = reg.replace('_', '')
                        defaults['reg'] = reg
                        if operator_ref in self.reg_operators:
                            condition |= Q(reg=reg)
                elif operator_ref in self.reg_operators:
                    reg = vehicle_ref.replace('_', '')
                    condition |= Q(reg=reg)
                elif operator_ref == 'WHIP':
                    code = vehicle_ref.replace('_', '')
                    condition |= Q(fleet_code=code)
        vehicles = vehicles.filter(condition)

        try:
            vehicle, created = vehicles.get_or_create(defaults)
        except Vehicle.MultipleObjectsReturned as e:
            print(e, operator, vehicle_ref)
            vehicle = vehicles.first()
            created = False

        self.vehicle_cache[cache_key] = vehicle.id
        return vehicle, created

    def get_service(self, operator, item, vehicle_operator_id):
        monitored_vehicle_journey = item['MonitoredVehicleJourney']

        line_ref = monitored_vehicle_journey.get('LineRef')
        if not line_ref:
            return

        destination_ref = monitored_vehicle_journey.get("DestinationRef")

        cache_key = f"{operator}:{vehicle_operator_id}:{line_ref}:{destination_ref}"
        if cache_key in self.service_cache:
            return self.service_cache[cache_key]

        services = Service.objects.filter(
            Exists(ServiceCode.objects.filter(service=OuterRef('id'), scheme__endswith=' SIRI', code=line_ref))
            | Q(line_name__iexact=line_ref),
            current=True
        )

        if type(operator) is Operator and operator.parent and destination_ref:
            services = services.filter(operator__parent=operator.parent)
            # we will use the destination ref to find out exactly which operator it is

        else:
            if type(operator) is Operator:
                condition = Q(operator=operator)
                if vehicle_operator_id != operator.id:
                    condition |= Q(operator=vehicle_operator_id)
                services = services.filter(condition)
            elif type(operator) is list:
                services = services.filter(operator__in=operator)

            try:
                return services.get()
            except Service.DoesNotExist:
                self.service_cache[cache_key] = None
                return
            except Service.MultipleObjectsReturned:
                pass

        if destination_ref:
            try:
                stops = StopPoint.objects.filter(service=OuterRef("pk"), locality__stoppoint=destination_ref)
                return services.filter(Exists(stops)).get()
            except Service.DoesNotExist:
                self.service_cache[cache_key] = None
                return
            except Service.MultipleObjectsReturned:
                try:
                    trips = Trip.objects.filter(route__service=OuterRef("pk"), destination=destination_ref)
                    return services.filter(Exists(trips)).get()
                except (Service.DoesNotExist, Service.MultipleObjectsReturned):
                    pass

        try:
            when = self.get_datetime(item)
            when = when.strftime('%a').lower()
            trips = Trip.objects.filter(**{f'calendar__{when}': True}, route__service=OuterRef("pk"))
            return services.filter(Exists(trips)).get()
        except (Service.DoesNotExist, Service.MultipleObjectsReturned):
            pass

        vehicle_journey_ref = monitored_vehicle_journey.get('VehicleJourneyRef')
        if vehicle_journey_ref and vehicle_journey_ref.isdigit():
            try:
                trips = Trip.objects.filter(route__service=OuterRef("pk"), ticket_machine_code=vehicle_journey_ref)
                return services.filter(Exists(trips)).get()
            except (Service.DoesNotExist, Service.MultipleObjectsReturned):
                pass

    def get_journey(self, item, vehicle):
        monitored_vehicle_journey = item['MonitoredVehicleJourney']

        vehicle_journey_ref = monitored_vehicle_journey.get('VehicleJourneyRef')
        if vehicle_journey_ref == 'UNKNOWN':
            vehicle_journey_ref = None

        route_name = monitored_vehicle_journey.get('LineRef') or ''

        origin_aimed_departure_time = monitored_vehicle_journey.get('OriginAimedDepartureTime')
        if origin_aimed_departure_time:
            origin_aimed_departure_time = parse_datetime(origin_aimed_departure_time)

        journey = None

        # do destination now, in case the route name is in the destination field
        destination = monitored_vehicle_journey.get('DestinationName') or ''
        if vehicle.operator_id == 'TGTC' and destination and not route_name:
            parts = destination.split()
            if parts[0].isdigit() or parts[0][:-1].isdigit():
                route_name = parts[0]
                destination = ' '.join(parts[1:])
                monitored_vehicle_journey['LineRef'] = route_name

        latest_location = vehicle.latest_location
        if latest_location:
            if origin_aimed_departure_time == latest_location.journey.datetime:
                journey = latest_location.journey
            elif vehicle_journey_ref:
                if vehicle_journey_ref != latest_location.journey.code and '_' in vehicle_journey_ref:
                    journey = vehicle.vehiclejourney_set.filter(route_name=route_name, code=vehicle_journey_ref).first()

            if not journey:
                journey = vehicle.vehiclejourney_set.filter(datetime=origin_aimed_departure_time).first()

        if not journey:
            journey = VehicleJourney(
                route_name=route_name,
                vehicle=vehicle,
                source=self.source,
                data=item,
                datetime=origin_aimed_departure_time,
                destination=destination
            )

        if vehicle_journey_ref:
            journey.code = vehicle_journey_ref

        if not journey.destination:
            destination_ref = monitored_vehicle_journey.get('DestinationRef')
            if destination_ref:
                try:
                    journey.destination = Locality.objects.get(stoppoint=destination_ref).name
                except Locality.DoesNotExist:
                    pass

        if (
            latest_location and latest_location.journey.service
            and latest_location.journey.code == journey.code
            and latest_location.journey.route_name == journey.route_name
        ):
            journey.service = latest_location.journey.service
        else:
            operator_ref = monitored_vehicle_journey["OperatorRef"]
            if operator_ref == "FWYO":
                if vehicle.operator_id == "HCTY":
                    operator = ["HCTY"]  # First West Yorkshire/Connexions
                else:
                    operator = ["FWYO", "FLDS"]
            elif operator_ref == "FHAL":
                operator = ["FHAL", "FHUD"]
            elif operator_ref == "FBRI" and vehicle.operator_id == "ABUS":
                operator = ["ABUS"]
            elif operator_ref == "RRAR":
                operator = ["RRAR", "FTVA"]  # Reading RailAir/First Berkshire
            elif operator_ref == "ROST":
                operator = ["ROST", "LNUD", "BPTR"]  # Rosso/Blackburn/Burnley
            else:
                operator = self.get_operator(operator_ref)
            journey.service = self.get_service(operator, item, vehicle.operator_id)

        return journey

    @staticmethod
    def create_vehicle_location(item):
        monitored_vehicle_journey = item['MonitoredVehicleJourney']
        location = monitored_vehicle_journey['VehicleLocation']
        latlong = Point(float(location['Longitude']), float(location['Latitude']))
        return VehicleLocation(
            latlong=latlong,
        )

    def has_changed(self, item):
        monitored_vehicle_journey = item['MonitoredVehicleJourney']
        key = f"{monitored_vehicle_journey['OperatorRef']}-{monitored_vehicle_journey['VehicleRef']}"
        if key in self.identifiers and self.identifiers[key] == item['RecordedAtTime']:
            return False
        self.identifiers[key] = item['RecordedAtTime']
        return True

    @async_to_sync
    async def send_items(self, items):
        await get_channel_layer().send('sirivm', {
            'type': 'sirivm',
            'items': items
        })

    def update(self):
        now = timezone.now()

        items = self.get_items()

        if items:
            chunk = []
            try:
                for item in items:
                    if self.has_changed(item):
                        chunk.append(item)
                    if len(chunk) == 50:
                        self.send_items(chunk)
                        chunk = []
                self.send_items(chunk)
            except ChannelFull:
                pass
        else:
            return 300  # no items - wait five minutes

        time_taken = timezone.now() - now
        print(time_taken)
        time_taken = time_taken.total_seconds()
        if time_taken < self.wait:
            return self.wait - time_taken
        return 0  # took longer than self.wait

    @staticmethod
    def items_from_response(response):
        return xmltodict.parse(
            response,
            dict_constructor=dict,  # override OrderedDict, cos dict is ordered in modern versions of Python
            force_list=['VehicleActivity']
        )

    def get_items(self):
        response = self.session.get(self.source.url, params=self.source.settings)
        if not response.ok:
            if 'datafeed' in self.source.url:
                print(response.content.decode())
            else:
                print(response)
            return

        if 'datafeed' in self.source.url:
            try:
                data = self.items_from_response(response.content)
            except ExpatError:
                print(response.content.decode())
                return
        else:
            try:
                with zipfile.ZipFile(io.BytesIO(response.content)) as archive:
                    assert archive.namelist() == ['siri.xml']
                    with archive.open('siri.xml') as open_file:
                        try:
                            data = self.items_from_response(open_file)
                        except ExpatError:
                            print(open_file.read())
                            return
            except zipfile.BadZipFile:
                print(response.content.decode())
                return

        self.source.datetime = parse_datetime(data['Siri']['ServiceDelivery']['ResponseTimestamp'])

        return data['Siri']['ServiceDelivery']['VehicleMonitoringDelivery']['VehicleActivity']
<|endoftext|>