SQLAlchemy在测试期间恢复auto_increment(pytest)

时间:2017-10-20 03:16:38

标签: python postgresql sqlalchemy pytest

我使用db_session

中的conftest.py灯具
@pytest.fixture(scope='function')
def db_session(db, request):
    """Creates a new database session for a test."""
    engine = create_engine(DefaultConfig.SQLALCHEMY_DATABASE_URI, connect_args={"options": "-c timezone=utc"})
    DbSession = sessionmaker(bind=engine)
    session = DbSession()

    connection = engine.connect()
    transaction = connection.begin()

    options = dict(bind=connection, binds={})
    session = db.create_scoped_session(options=options)

    db.session = session

    yield session

    transaction.rollback()
    connection.close()
    session.remove()

在我的测试脚本test_project.py

import pytest
from package.model import Project

def test_create_project(db_session):
    project = Project(project_name="Test_Project")

    db_session.add(project)
    db_session.commit()

    assert project.project_id > 0
    assert project.project_id == 65 # supposed to be the next auto-increment value
    assert db_session.query(Project).filter_by(project_name="Test_Project").first() != None

每次运行测试时,自动递增的主键都会增加。即使我使用db_session fixture,主键序列也不会回滚。

我还尝试在测试之间调用db_session.begin_nested()db_session.rollback(),但主键仍然增加

def test_create_project(db_session):
    db_session.begin_nested()
    project = Project(project_name="Test_Project")

    db_session.add(project)
    db_session.commit()

    assert project.project_id > 0
    assert project.project_id == 65
    assert db_session.query(Project).filter_by(project_name="Test_Project").first() != None
    db_session.rollback()

如何在测试后恢复自动增量值?

环境

  • 数据库:PostgreSQL 9.6.5
  • 烧瓶0.12.2
  • SQLAlchemy 1.1.12

1 个答案:

答案 0 :(得分:1)

使用sqlalchemy无法做到这一点。但是,你可以改变&amp;在所需的起始点重新启动序列。默认情况下,postgresql会将序列命名为<table_name>_<column_name>_seq。检查您的架构以找出序列名称。需要改变序列。

这是一个小脚本来说明这一点:

CREATE TABLE test(col1 SERIAL PRIMARY KEY, val TEXT);
INSERT INTO test (val) VALUES ('abc'), ('def'), ('ghi');
SELECT * FROM test;
-- output
1   abc
2   def
3   ghi

DELETE FROM test WHERE col1 = 3;
SELECT * FROM test;
-- output
1   abc
2   def

-- now alter the sequence so that the next insert gets value 3 in col1
ALTER SEQUENCE test_col1_seq START 3 RESTART;
INSERT INTO test (val) VALUES ('xyz');
SELECT * FROM test;
-- output
1   abc
2   def
3   xyz