jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

fix: Use AsyncSession for user management (#4491)

Browse source 
* Use AsyncSession for user management

* Simplify check_key

* Don't trigger blockbuster on settings service initialize

* Fix mypy

* Fix api key update_total_uses

* Fix auto-login

* Revert making CustomComponent.list_key_names async
This commit is contained in:
Christophe Bornet 2024-11-16 02:09:33 +01:00 • committed by GitHub
commit 6573ca14cc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
24 changed files with 430 additions and 339 deletions
Download patch file Download diff file Expand all files Collapse all files

121
src/backend/tests/unit/services/variable/test_service.py
View file

@ -1,6 +1,6 @@
from datetime import datetime
from unittest.mock import patch
from uuid import uuid4
from uuid import UUID, uuid4
import pytest
from langflow.services.database.models.variable.model import VariableUpdate
@ -8,7 +8,9 @@ from langflow.services.deps import get_settings_service
from langflow.services.settings.constants import VARIABLES_TO_GET_FROM_ENVIRONMENT
from langflow.services.variable.constants import CREDENTIAL_TYPE, GENERIC_TYPE
from langflow.services.variable.service import DatabaseVariableService
from sqlmodel import Session, SQLModel, create_engine
from sqlalchemy.ext.asyncio import create_async_engine
from sqlmodel import Session, SQLModel
from sqlmodel.ext.asyncio.session import AsyncSession
@pytest.fixture
@ -18,114 +20,125 @@ def service():
@pytest.fixture
def session():
engine = create_engine("sqlite:///:memory:")
SQLModel.metadata.create_all(engine)
with Session(engine) as session:
async def session():
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
async with engine.begin() as conn:
await conn.run_sync(SQLModel.metadata.create_all)
async with AsyncSession(engine) as session:
yield session
def test_initialize_user_variables__create_and_update(service, session):
def _get_variable(
session: Session,
service,
user_id: UUID | str,
name: str,
field: str,
):
return service.get_variable(user_id, name, field, session=session)
async def test_initialize_user_variables__create_and_update(service, session: AsyncSession):
user_id = uuid4()
field = ""
good_vars = {k: f"value{i}" for i, k in enumerate(VARIABLES_TO_GET_FROM_ENVIRONMENT)}
bad_vars = {"VAR1": "value1", "VAR2": "value2", "VAR3": "value3"}
env_vars = {**good_vars, **bad_vars}
service.create_variable(user_id, "OPENAI_API_KEY", "outdate", session=session)
await service.create_variable(user_id, "OPENAI_API_KEY", "outdate", session=session)
env_vars["OPENAI_API_KEY"] = "updated_value"
with patch.dict("os.environ", env_vars, clear=True):
service.initialize_user_variables(user_id=user_id, session=session)
await service.initialize_user_variables(user_id=user_id, session=session)
variables = service.list_variables(user_id, session=session)
variables = await service.list_variables(user_id, session=session)
for name in variables:
value = service.get_variable(user_id, name, field, session=session)
value = await session.run_sync(_get_variable, service, user_id, name, field)
assert value == env_vars[name]
assert all(i in variables for i in good_vars)
assert all(i not in variables for i in bad_vars)
def test_initialize_user_variables__not_found_variable(service, session):
async def test_initialize_user_variables__not_found_variable(service, session: AsyncSession):
with patch("langflow.services.variable.service.DatabaseVariableService.create_variable") as m:
m.side_effect = Exception()
service.initialize_user_variables(uuid4(), session=session)
await service.initialize_user_variables(uuid4(), session=session)
assert True
def test_initialize_user_variables__skipping_environment_variable_storage(service, session):
async def test_initialize_user_variables__skipping_environment_variable_storage(service, session: AsyncSession):
service.settings_service.settings.store_environment_variables = False
service.initialize_user_variables(uuid4(), session=session)
await service.initialize_user_variables(uuid4(), session=session)
assert True
def test_get_variable(service, session):
async def test_get_variable(service, session: AsyncSession):
user_id = uuid4()
name = "name"
value = "value"
field = ""
service.create_variable(user_id, name, value, session=session)
await service.create_variable(user_id, name, value, session=session)
result = service.get_variable(user_id, name, field, session=session)
result = await session.run_sync(_get_variable, service, user_id, name, field)
assert result == value
def test_get_variable__valueerror(service, session):
async def test_get_variable__valueerror(service, session: AsyncSession):
user_id = uuid4()
name = "name"
field = ""
with pytest.raises(ValueError, match=f"{name} variable not found."):
service.get_variable(user_id, name, field, session)
await session.run_sync(_get_variable, service, user_id, name, field)
def test_get_variable__typeerror(service, session):
async def test_get_variable__typeerror(service, session: AsyncSession):
user_id = uuid4()
name = "name"
value = "value"
field = "session_id"
_type = CREDENTIAL_TYPE
service.create_variable(user_id, name, value, _type=_type, session=session)
await service.create_variable(user_id, name, value, _type=_type, session=session)
with pytest.raises(TypeError) as exc:
service.get_variable(user_id, name, field, session)
await session.run_sync(_get_variable, service, user_id, name, field)
assert name in str(exc.value)
assert "purpose is to prevent the exposure of value" in str(exc.value)
def test_list_variables(service, session):
async def test_list_variables(service, session: AsyncSession):
user_id = uuid4()
names = ["name1", "name2", "name3"]
value = "value"
for name in names:
service.create_variable(user_id, name, value, session=session)
await service.create_variable(user_id, name, value, session=session)
result = service.list_variables(user_id, session=session)
result = await service.list_variables(user_id, session=session)
assert all(name in result for name in names)
def test_list_variables__empty(service, session):
result = service.list_variables(uuid4(), session=session)
async def test_list_variables__empty(service, session: AsyncSession):
result = await service.list_variables(uuid4(), session=session)
assert not result
assert isinstance(result, list)
def test_update_variable(service, session):
async def test_update_variable(service, session: AsyncSession):
user_id = uuid4()
name = "name"
old_value = "old_value"
new_value = "new_value"
field = ""
service.create_variable(user_id, name, old_value, session=session)
await service.create_variable(user_id, name, old_value, session=session)
old_recovered = service.get_variable(user_id, name, field, session=session)
result = service.update_variable(user_id, name, new_value, session=session)
new_recovered = service.get_variable(user_id, name, field, session=session)
old_recovered = await session.run_sync(_get_variable, service, user_id, name, field)
result = await service.update_variable(user_id, name, new_value, session=session)
new_recovered = await session.run_sync(_get_variable, service, user_id, name, field)
assert old_value == old_recovered
assert new_value == new_recovered
@ -139,26 +152,26 @@ def test_update_variable(service, session):
assert isinstance(result.updated_at, datetime)
def test_update_variable__valueerror(service, session):
async def test_update_variable__valueerror(service, session: AsyncSession):
user_id = uuid4()
name = "name"
value = "value"
with pytest.raises(ValueError, match=f"{name} variable not found."):
service.update_variable(user_id, name, value, session=session)
await service.update_variable(user_id, name, value, session=session)
def test_update_variable_fields(service, session):
async def test_update_variable_fields(service, session: AsyncSession):
user_id = uuid4()
new_name = new_value = "donkey"
variable = service.create_variable(user_id, "old_name", "old_value", session=session)
variable = await service.create_variable(user_id, "old_name", "old_value", session=session)
saved = variable.model_dump()
variable = VariableUpdate(**saved)
variable.name = new_name
variable.value = new_value
variable.default_fields = ["new_field"]
result = service.update_variable_fields(
result = await service.update_variable_fields(
user_id=user_id,
variable_id=saved.get("id"),
variable=variable,
@ -177,58 +190,58 @@ def test_update_variable_fields(service, session):
assert saved.get("updated_at") != result.updated_at
def test_delete_variable(service, session):
async def test_delete_variable(service, session: AsyncSession):
user_id = uuid4()
name = "name"
value = "value"
field = ""
service.create_variable(user_id, name, value, session=session)
recovered = service.get_variable(user_id, name, field, session=session)
service.delete_variable(user_id, name, session=session)
await service.create_variable(user_id, name, value, session=session)
recovered = await session.run_sync(_get_variable, service, user_id, name, field)
await service.delete_variable(user_id, name, session=session)
with pytest.raises(ValueError, match=f"{name} variable not found."):
service.get_variable(user_id, name, field, session)
await session.run_sync(_get_variable, service, user_id, name, field)
assert recovered == value
def test_delete_variable__valueerror(service, session):
async def test_delete_variable__valueerror(service, session: AsyncSession):
user_id = uuid4()
name = "name"
with pytest.raises(ValueError, match=f"{name} variable not found."):
service.delete_variable(user_id, name, session=session)
await service.delete_variable(user_id, name, session=session)
def test_delete_variable_by_id(service, session):
async def test_delete_variable_by_id(service, session: AsyncSession):
user_id = uuid4()
name = "name"
value = "value"
field = "field"
saved = service.create_variable(user_id, name, value, session=session)
recovered = service.get_variable(user_id, name, field, session=session)
service.delete_variable_by_id(user_id, saved.id, session=session)
saved = await service.create_variable(user_id, name, value, session=session)
recovered = await session.run_sync(_get_variable, service, user_id, name, field)
await service.delete_variable_by_id(user_id, saved.id, session=session)
with pytest.raises(ValueError, match=f"{name} variable not found."):
service.get_variable(user_id, name, field, session)
await session.run_sync(_get_variable, service, user_id, name, field)
assert recovered == value
def test_delete_variable_by_id__valueerror(service, session):
async def test_delete_variable_by_id__valueerror(service, session: AsyncSession):
user_id = uuid4()
variable_id = uuid4()
with pytest.raises(ValueError, match=f"{variable_id} variable not found."):
service.delete_variable_by_id(user_id, variable_id, session=session)
await service.delete_variable_by_id(user_id, variable_id, session=session)
def test_create_variable(service, session):
async def test_create_variable(service, session: AsyncSession):
user_id = uuid4()
name = "name"
value = "value"
result = service.create_variable(user_id, name, value, session=session)
result = await service.create_variable(user_id, name, value, session=session)
assert result.user_id == user_id
assert result.name == name