ficapy
12/14/2016 - 7:10 AM

sqlalchemy+pytest

sqlalchemy+pytest

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Ficapy
# Create: '12/14/16'

import pytest
from sqlalchemy import create_engine, Column, BIGINT
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.exc import IntegrityError

sqlalchemy_db = 'postgresql+psycopg2://user:pwd@127.0.0.1/db'

Base = declarative_base()


class Demo(Base):
    __tablename__ = 'demo'
    id = Column(BIGINT, primary_key=True)
    task = Column(BIGINT)


uline_engine = create_engine(sqlalchemy_db, pool_recycle=3600, echo=False)

Base.metadata.bind = uline_engine
Base.metadata.create_all(checkfirst=True)
db_Session = sessionmaker()

session = db_Session()


@pytest.fixture(autouse=True)
def trade_session(request):
    # 新建一个单独的数据库连接
    conn = uline_engine.connect()
    Base.metadata.bind = conn

    # 开始一个非ORM事务
    trans = conn.begin()

    # 绑定这个连接到db_Session
    db_Session.configure(bind=conn)

    def tearDown():
        session.close()
        trans.rollback()
        conn.close()

    request.addfinalizer(tearDown)


def insert_data(id, task):
    session.add(Demo(id=id, task=task))
    session.commit()


def test_insert_data(trade_session):
    insert_data(20, 2)

    assert session.query(Demo).filter_by(id=20).first() is not None
    with pytest.raises(IntegrityError):
        insert_data(20, 2)