import enum
from sqlalchemy import (
    create_engine,
    MetaData,
    Table,
    Column,
    Enum,
    bindparam,
    select,
    func,
)
from sqlalchemy_utils import (
    database_exists,
    create_database,
)


class TestEnum(enum.Enum):
    FOO = 0
    BAR = 1
    BAZ = 2


URL = 'sqlite:////tmp/db.db'


if not database_exists(URL):
    create_database(URL)


ENGINE = create_engine(URL)
METADATA = MetaData(bind=ENGINE)


TABLE = Table(
    'test',
    METADATA,
    Column('test', Enum(TestEnum), nullable=True),
)

if ENGINE.dialect.has_table(ENGINE, 'test'):
    METADATA.drop_all()
METADATA.create_all()


with ENGINE.connect() as c:
    c.execute(TABLE.insert([
        dict(test=TestEnum.FOO),
        dict(test=TestEnum.BAR),
        dict(test=TestEnum.BAZ),
    ]))


QUERY = select([
    func.count(),
]).where(
    TABLE.c.test.in_(bindparam('test', expanding=True)),
)


with ENGINE.connect() as c:
    result = c.execute(QUERY.params(
        test=[TestEnum.FOO.name, TestEnum.BAR.name],
    )).scalar()
    assert result == 2
    result = c.execute(QUERY.params(
        test=[TestEnum.FOO, TestEnum.BAR],
    )).scalar()
    assert result == 2