<commit_before>import os
import shutil
import tempfile
import unittest

import sqlalchemy as sa
from sqlalchemy.engine.base import Engine
import sqlalchemy.ext.declarative as declarative

import sqlahelper

class DBInfo(object):
    def __init__(self, dir, filename):
        self.file = os.path.join(dir, filename)
        self.url = "sqlite:///" + self.file

class SQLAHelperTestCase(unittest.TestCase):
    def setUp(self):
        self.dir = tempfile.mkdtemp()
        self.db1 = DBInfo(self.dir, "db1.sqlite")
        self.db2 = DBInfo(self.dir, "db2.sqlite")
        self.db3 = DBInfo(self.dir, "db3.sqlite")

    def tearDown(self):
        sqlahelper.reset()
        shutil.rmtree(self.dir, True)

    if not hasattr(unittest.TestCase, "assertIsInstance"): # pragma: no cover
        def assertIsInstance(self, obj, classes):
            if not isinstance(obj, classes):
                typ = type(obj)
                if isinstance(classes, (list, tuple)):
                    classes_str = ", ".join(x.__name__ for x in classes)
                    classes_str = "[%s]" % classes_str
                else:
                    classes_str = classes.__name__
                msg = "%s is not an instance of %s" % (typ, classes_str)
                raise AssertionError(msg)

        def assertIs(self, a, b):
            if a is not b:
                raise AssertionError("%r is not %r" % (a, b))


class TestAddEngine(SQLAHelperTestCase):
    def test_one_engine(self):
        e = sa.create_engine(self.db1.url)
        sqlahelper.add_engine(e)
        retrieved = sqlahelper.get_engine()
        self.assertIs(retrieved, e)

    def test_multiple_engines(self):
        default = sa.create_engine(self.db1.url)
        stats = sa.create_engine(self.db2.url)
        sqlahelper.add_engine(default)
        sqlahelper.add_engine(stats, "stats")
        # Can we retrieve the engines?
        self.assertIs(sqlahelper.get_engine(), default)
        self.assertIs(sqlahelper.get_engine("default"), default)
        self.assertIs(sqlahelper.get_engine("stats"), stats)
        # Are the session binding and base binding set correctly?
        self.assertIs(sqlahelper.get_session().bind, default)
        self.assertIs(sqlahelper.get_base().metadata.bind, default)

    def test_multiple_engines_without_default(self):
        db1 = sa.create_engine(self.db1.url)
        db2 = sa.create_engine(self.db2.url)
        sqlahelper.add_engine(db1, "db1")
        sqlahelper.add_engine(db2, "db2")
        # Can we retrieve the engines?
        self.assertIs(sqlahelper.get_engine("db1"), db1)
        self.assertIs(sqlahelper.get_engine("db2"), db2)
        # There should be no default engine
        self.assertIsNone(sqlahelper.get_session().bind)
        self.assertIsNone(sqlahelper.get_base().metadata.bind)
        self.assertIsNone(sqlahelper.get_engine())


class TestDeclarativeBase(SQLAHelperTestCase):
    def test1(self):
        import transaction
        Base = sqlahelper.get_base()
        class Person(Base):
            __tablename__ = "people"
            id = sa.Column(sa.Integer, primary_key=True)
            first_name = sa.Column(sa.Unicode(100), nullable=False)
            last_name = sa.Column(sa.Unicode(100), nullable=False)
        engine = sa.create_engine(self.db1.url)
        sqlahelper.add_engine(engine)
        Base.metadata.create_all()
        fred = Person(id=1, first_name=u"Fred", last_name=u"Flintstone")
        wilma = Person(id=2, first_name=u"Wilma", last_name=u"Flintstone")
        barney = Person(id=3, first_name=u"Barney", last_name=u"Rubble")
        betty = Person(id=4, first_name=u"Betty", last_name=u"Rubble")<commit_msg>add test case<commit_after>import os
import shutil
import tempfile
import unittest

import sqlalchemy as sa
from sqlalchemy.engine.base import Engine
import sqlalchemy.ext.declarative as declarative

import sqlahelper

class DBInfo(object):
    def __init__(self, dir, filename):
        self.file = os.path.join(dir, filename)
        self.url = "sqlite:///" + self.file

class SQLAHelperTestCase(unittest.TestCase):
    def setUp(self):
        self.dir = tempfile.mkdtemp()
        self.db1 = DBInfo(self.dir, "db1.sqlite")
        self.db2 = DBInfo(self.dir, "db2.sqlite")
        self.db3 = DBInfo(self.dir, "db3.sqlite")

    def tearDown(self):
        sqlahelper.reset()
        shutil.rmtree(self.dir, True)

    if not hasattr(unittest.TestCase, "assertIsInstance"): # pragma: no cover
        def assertIsInstance(self, obj, classes):
            if not isinstance(obj, classes):
                typ = type(obj)
                if isinstance(classes, (list, tuple)):
                    classes_str = ", ".join(x.__name__ for x in classes)
                    classes_str = "[%s]" % classes_str
                else:
                    classes_str = classes.__name__
                msg = "%s is not an instance of %s" % (typ, classes_str)
                raise AssertionError(msg)

        def assertIs(self, a, b):
            if a is not b:
                raise AssertionError("%r is not %r" % (a, b))


class TestAddEngine(SQLAHelperTestCase):
    def test_one_engine(self):
        e = sa.create_engine(self.db1.url)
        sqlahelper.add_engine(e)
        retrieved = sqlahelper.get_engine()
        self.assertIs(retrieved, e)

    def test_multiple_engines(self):
        default = sa.create_engine(self.db1.url)
        stats = sa.create_engine(self.db2.url)
        sqlahelper.add_engine(default)
        sqlahelper.add_engine(stats, "stats")
        # Can we retrieve the engines?
        self.assertIs(sqlahelper.get_engine(), default)
        self.assertIs(sqlahelper.get_engine("default"), default)
        self.assertIs(sqlahelper.get_engine("stats"), stats)
        # Are the session binding and base binding set correctly?
        self.assertIs(sqlahelper.get_session().bind, default)
        self.assertIs(sqlahelper.get_base().metadata.bind, default)

    def test_multiple_engines_without_default(self):
        db1 = sa.create_engine(self.db1.url)
        db2 = sa.create_engine(self.db2.url)
        sqlahelper.add_engine(db1, "db1")
        sqlahelper.add_engine(db2, "db2")
        # Can we retrieve the engines?
        self.assertIs(sqlahelper.get_engine("db1"), db1)
        self.assertIs(sqlahelper.get_engine("db2"), db2)
        # There should be no default engine
        self.assertIsNone(sqlahelper.get_session().bind)
        self.assertIsNone(sqlahelper.get_base().metadata.bind)
        self.assertIsNone(sqlahelper.get_engine())

    def test_add_engine_twice(self):
        db1 = sa.create_engine(self.db1.url)
        db2 = sa.create_engine(self.db2.url)
        sqlahelper.add_engine(db1)
        self.assertIs(sqlahelper.get_session().bind, db1)
        sqlahelper.add_engine(db2)
        self.assertIs(sqlahelper.get_session().bind, db2)
        self.assertIs(sqlahelper.get_session().bind, sqlahelper.sessions.default.registry.registry.value.bind)



class TestDeclarativeBase(SQLAHelperTestCase):
    def test1(self):
        import transaction
        Base = sqlahelper.get_base()
        class Person(Base):
            __tablename__ = "people"
            id = sa.Column(sa.Integer, primary_key=True)
            first_name = sa.Column(sa.Unicode(100), nullable=False)
            last_name = sa.Column(sa.Unicode(100), nullable=False)
        engine = sa.create_engine(self.db1.url)
        sqlahelper.add_engine(engine)
        Base.metadata.create_all()
        fred = Person(id=1, first_name=u"Fred", last_name=u"Flintstone")
        wilma = Person(id=2, first_name=u"Wilma", last_name=u"Flintstone")
        barney = Person(id=3, first_name=u"Barney", last_name=u"Rubble")
        betty = Person(id=4, first_name=u"Betty", last_name=u"Rubble")
{'fraction_non_alphanumeric': 0.11409137577002053, 'fraction_numerical': 0.009496919917864477, 'mean_word_length': 2.7144899904671114, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '430653', 'n_tokens_mistral': 2533, 'n_tokens_neox': 2480, 'n_words': 471}