from unittest import TestCase, mock
from sqlalchemy import Column, BigInteger, String

from ._common import db, mysql_url, mysql_url_1, mysql_url_2


class Foo(db.Model):
    __bind_key__ = 'foo'
    __tablename__ = 'foo'

    id = Column(BigInteger, primary_key=True)
    foo = Column(String(64))

    def __init__(self, foo):
        self.foo = foo


class Bar(db.Model):
    __bind_key__ = 'bar'
    __tablename__ = 'bar'

    id = Column(BigInteger, primary_key=True)
    bar = Column(String(64))

    def __init__(self, bar):
        self.bar = bar


class MultipleDatabasesTestCase(TestCase):
    def setUp(self, *args, **kwargs):
        super(MultipleDatabasesTestCase, self).setUp(*args, **kwargs)

        db.configure(
            uri=mysql_url, binds={'foo': mysql_url_1, 'bar': mysql_url_2}
        )

        self._application = mock.Mock()
        self._application.settings = {'db': db}

        db.create_all()

    def tearDown(self, *args, **kwargs):
        db.drop_all()

        super(MultipleDatabasesTestCase, self).tearDown(*args, **kwargs)

    def test_add_objects(self):
        session = db.sessionmaker()

        session.add(Foo('foo'))
        session.commit()

        session.add(Bar('first bar'))
        session.add(Bar('second bar'))
        session.commit()

        foo_count = (
            db.get_engine('foo')
            .execute("SELECT COUNT(*) FROM foo")
            .fetchone()[0]
        )
        bar_count = (
            db.get_engine('bar')
            .execute("SELECT COUNT(*) FROM bar")
            .fetchone()[0]
        )

        session.close()

        self.assertEqual(foo_count, 1)
        self.assertEqual(bar_count, 2)

@@ -1,12 +1,40 @@
+from http import HTTPStatus%0A
 from unittes
@@ -57,16 +57,17 @@
 e, mock%0A
+%0A
 from sql
@@ -107,16 +107,159 @@
 , String
+%0Afrom tornado.testing import AsyncHTTPTestCase%0Afrom tornado.web import Application, RequestHandler%0A%0Afrom tornado_sqlalchemy import SessionMixin
 %0A%0Afrom .
@@ -1848,8 +1848,1452 @@
 unt, 2)%0A
+%0A%0Aclass RequestHandlersTestCase(AsyncHTTPTestCase, TestCase):%0A    def __init__(self, *args, **kwargs):%0A        super(RequestHandlersTestCase, self).__init__(*args, **kwargs)%0A%0A        db.configure(%0A            uri=mysql_url, binds=%7B'foo': mysql_url_1, 'bar': mysql_url_2%7D%0A        )%0A%0A        class IndexRequestHandler(SessionMixin, RequestHandler):%0A            def get(h_self):%0A                with h_self.make_session() as session:%0A                    session.add(Foo('first-foo'))%0A%0A                    session.add(Bar('first-bar'))%0A                    session.add(Bar('second-bar'))%0A%0A                    session.flush()%0A%0A                    foo_count = session.query(Foo).count()%0A                    bar_count = session.query(Bar).count()%0A%0A                h_self.write('%7B%7D, %7B%7D'.format(foo_count, bar_count))%0A%0A        handlers = ((r'/', IndexRequestHandler),)%0A%0A        self._application = Application(handlers, db=db)%0A%0A    def setUp(self, *args, **kwargs):%0A        super(RequestHandlersTestCase, self).setUp(*args, **kwargs)%0A%0A        db.create_all()%0A%0A    def tearDown(self, *args, **kwargs):%0A        db.drop_all()%0A%0A        super(RequestHandlersTestCase, self).tearDown(*args, **kwargs)%0A%0A    def get_app(self):%0A        return self._application%0A%0A    def test_count(self):%0A        response = self.fetch('/', method='GET')%0A%0A        self.assertEqual(response.code, HTTPStatus.OK.value)%0A        self.assertEqual(response.body.decode('utf-8'), '1, 2')%0A

add request handler test cases for multiple databases