diff --git a/tests/integration/db_migrations/test_autoincrement_primary_keys.py b/tests/integration/db_migrations/test_autoincrement_primary_keys.py index 85586f75b3..bd7d68ffbd 100644 --- a/tests/integration/db_migrations/test_autoincrement_primary_keys.py +++ b/tests/integration/db_migrations/test_autoincrement_primary_keys.py @@ -2,10 +2,7 @@ from alembic.config import Config from phoenix.db import models -from sqlalchemy import ( - Engine, - select, -) +from sqlalchemy import Engine, select from sqlalchemy.orm import sessionmaker from . import _up @@ -32,11 +29,13 @@ def test_prompt_versions( prompt_version = models.PromptVersion(**values) session.add(prompt_version) assert (id_ := prompt_version.id) is not None + with db.begin() as session: + assert session.scalar(select(models.PromptVersion.id).filter_by(id=id_)) is not None with db.begin() as session: session.delete(prompt_version) with db.begin() as session: assert session.scalar(select(models.PromptVersion.id).filter_by(id=id_)) is None with db.begin() as session: - prompt_version = models.PromptVersion(**values) - session.add(prompt_version) - assert prompt_version.id > id_ + new_prompt_version = models.PromptVersion(**values) + session.add(new_prompt_version) + assert new_prompt_version.id > id_