🐛 fix(base.py): use db_manager.engine instead of DatabaseManager.engine to access the database engine

🐛 fix(conftest.py): add TYPE_CHECKING import to fix type hinting error
🐛 fix(conftest.py): pass db_manager to blank_session_getter fixture to fix session creation error
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-08-06 12:15:29 -03:00
commit a9db2da6bf
2 changed files with 8 additions and 5 deletions

View file

@ -70,7 +70,7 @@ class DatabaseManager(Service):
@contextmanager @contextmanager
def session_getter(db_manager: DatabaseManager): def session_getter(db_manager: DatabaseManager):
try: try:
session = Session(DatabaseManager.engine) session = Session(db_manager.engine)
yield session yield session
except Exception as e: except Exception as e:
print("Session rollback because of exception:", e) print("Session rollback because of exception:", e)

View file

@ -1,7 +1,7 @@
from contextlib import contextmanager from contextlib import contextmanager
import json import json
from pathlib import Path from pathlib import Path
from typing import AsyncGenerator from typing import AsyncGenerator, TYPE_CHECKING
from langflow.api.v1.flows import get_session from langflow.api.v1.flows import get_session
from langflow.graph.graph.base import Graph from langflow.graph.graph.base import Graph
@ -11,6 +11,9 @@ from httpx import AsyncClient
from sqlmodel import SQLModel, Session, create_engine from sqlmodel import SQLModel, Session, create_engine
from sqlmodel.pool import StaticPool from sqlmodel.pool import StaticPool
if TYPE_CHECKING:
from langflow.services.database.base import DatabaseManager
def pytest_configure(): def pytest_configure():
pytest.BASIC_EXAMPLE_PATH = ( pytest.BASIC_EXAMPLE_PATH = (
@ -134,15 +137,15 @@ def client_fixture(session: Session):
# create a fixture for session_getter above # create a fixture for session_getter above
@pytest.fixture(name="session_getter") @pytest.fixture(name="session_getter")
def session_getter_fixture(): def session_getter_fixture(client):
engine = create_engine( engine = create_engine(
"sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool "sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool
) )
SQLModel.metadata.create_all(engine) SQLModel.metadata.create_all(engine)
@contextmanager @contextmanager
def blank_session_getter(): def blank_session_getter(db_manager: "DatabaseManager"):
with Session(engine) as session: with Session(db_manager.engine) as session:
yield session yield session
yield blank_session_getter yield blank_session_getter