community/SQLDatabase: Generalize and trim software tests (#16659)

- **Description:** Improve test cases for `SQLDatabase` adapter
component, see
[suggestion](https://github.com/langchain-ai/langchain/pull/16655#pullrequestreview-1846749474).
  - **Depends on:** GH-16655
  - **Addressed to:** @baskaryan, @cbornet, @eyurtsev

_Remark: This PR is stacked upon GH-16655, so that one will need to go
in first._

Edit: Thank you for bringing in GH-17191, @eyurtsev. This is a little
aftermath, improving/streamlining the corresponding test cases.
pull/17441/head
Andreas Motl 8 months ago committed by GitHub
parent 1987f905ed
commit 1fdd9bd980
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -10,7 +10,6 @@ from sqlalchemy import (
String, String,
Table, Table,
Text, Text,
create_engine,
insert, insert,
select, select,
) )
@ -35,11 +34,19 @@ company = Table(
) )
def test_table_info() -> None: @pytest.fixture
"""Test that table info is constructed properly.""" def engine() -> sa.Engine:
engine = create_engine("sqlite:///:memory:") return sa.create_engine("sqlite:///:memory:")
@pytest.fixture
def db(engine: sa.Engine) -> SQLDatabase:
metadata_obj.create_all(engine) metadata_obj.create_all(engine)
db = SQLDatabase(engine) return SQLDatabase(engine)
def test_table_info(db: SQLDatabase) -> None:
"""Test that table info is constructed properly."""
output = db.table_info output = db.table_info
expected_output = """ expected_output = """
CREATE TABLE user ( CREATE TABLE user (
@ -68,20 +75,19 @@ def test_table_info() -> None:
assert sorted(" ".join(output.split())) == sorted(" ".join(expected_output.split())) assert sorted(" ".join(output.split())) == sorted(" ".join(expected_output.split()))
def test_table_info_w_sample_rows() -> None: def test_table_info_w_sample_rows(db: SQLDatabase) -> None:
"""Test that table info is constructed properly.""" """Test that table info is constructed properly."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine) # Provision.
values = [ values = [
{"user_id": 13, "user_name": "Harrison", "user_bio": "bio"}, {"user_id": 13, "user_name": "Harrison", "user_bio": "bio"},
{"user_id": 14, "user_name": "Chase", "user_bio": "bio"}, {"user_id": 14, "user_name": "Chase", "user_bio": "bio"},
] ]
stmt = insert(user).values(values) stmt = insert(user).values(values)
with engine.begin() as conn: db._execute(stmt)
conn.execute(stmt)
db = SQLDatabase(engine, sample_rows_in_table_info=2)
# Query and verify.
db = SQLDatabase(db._engine, sample_rows_in_table_info=2)
output = db.table_info output = db.table_info
expected_output = """ expected_output = """
@ -112,16 +118,16 @@ def test_table_info_w_sample_rows() -> None:
assert sorted(output.split()) == sorted(expected_output.split()) assert sorted(output.split()) == sorted(expected_output.split())
def test_sql_database_run_fetch_all() -> None: def test_sql_database_run_fetch_all(db: SQLDatabase) -> None:
"""Verify running SQL expressions returning results as strings.""" """Verify running SQL expressions returning results as strings."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine) # Provision.
stmt = insert(user).values( stmt = insert(user).values(
user_id=13, user_name="Harrison", user_bio="That is my Bio " * 24 user_id=13, user_name="Harrison", user_bio="That is my Bio " * 24
) )
with engine.begin() as conn: db._execute(stmt)
conn.execute(stmt)
db = SQLDatabase(engine) # Query and verify.
command = "select user_id, user_name, user_bio from user where user_id = 13" command = "select user_id, user_name, user_bio from user where user_id = 13"
partial_output = db.run(command) partial_output = db.run(command)
user_bio = "That is my Bio " * 19 + "That is my..." user_bio = "That is my Bio " * 19 + "That is my..."
@ -135,60 +141,57 @@ def test_sql_database_run_fetch_all() -> None:
assert full_output == expected_full_output assert full_output == expected_full_output
def test_sql_database_run_fetch_result() -> None: def test_sql_database_run_fetch_result(db: SQLDatabase) -> None:
"""Verify running SQL expressions returning results as SQLAlchemy `Result` instances.""" """Verify running SQL expressions returning results as SQLAlchemy `Result` instances."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine) # Provision.
stmt = insert(user).values(user_id=17, user_name="hwchase") stmt = insert(user).values(user_id=17, user_name="hwchase")
with engine.begin() as conn: db._execute(stmt)
conn.execute(stmt)
db = SQLDatabase(engine)
command = "select user_id, user_name, user_bio from user where user_id = 17"
# Query and verify.
command = "select user_id, user_name, user_bio from user where user_id = 17"
result = db.run(command, fetch="cursor", include_columns=True) result = db.run(command, fetch="cursor", include_columns=True)
expected = [{"user_id": 17, "user_name": "hwchase", "user_bio": None}] expected = [{"user_id": 17, "user_name": "hwchase", "user_bio": None}]
assert isinstance(result, Result) assert isinstance(result, Result)
assert result.mappings().fetchall() == expected assert result.mappings().fetchall() == expected
def test_sql_database_run_with_parameters() -> None: def test_sql_database_run_with_parameters(db: SQLDatabase) -> None:
"""Verify running SQL expressions with query parameters.""" """Verify running SQL expressions with query parameters."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine) # Provision.
stmt = insert(user).values(user_id=17, user_name="hwchase") stmt = insert(user).values(user_id=17, user_name="hwchase")
with engine.begin() as conn: db._execute(stmt)
conn.execute(stmt)
db = SQLDatabase(engine)
command = "select user_id, user_name, user_bio from user where user_id = :user_id"
# Query and verify.
command = "select user_id, user_name, user_bio from user where user_id = :user_id"
full_output = db.run(command, parameters={"user_id": 17}, include_columns=True) full_output = db.run(command, parameters={"user_id": 17}, include_columns=True)
expected_full_output = "[{'user_id': 17, 'user_name': 'hwchase', 'user_bio': None}]" expected_full_output = "[{'user_id': 17, 'user_name': 'hwchase', 'user_bio': None}]"
assert full_output == expected_full_output assert full_output == expected_full_output
def test_sql_database_run_sqlalchemy_selectable() -> None: def test_sql_database_run_sqlalchemy_selectable(db: SQLDatabase) -> None:
"""Verify running SQL expressions using SQLAlchemy selectable.""" """Verify running SQL expressions using SQLAlchemy selectable."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine) # Provision.
stmt = insert(user).values(user_id=17, user_name="hwchase") stmt = insert(user).values(user_id=17, user_name="hwchase")
with engine.begin() as conn: db._execute(stmt)
conn.execute(stmt)
db = SQLDatabase(engine)
command = select(user).where(user.c.user_id == 17)
# Query and verify.
command = select(user).where(user.c.user_id == 17)
full_output = db.run(command, include_columns=True) full_output = db.run(command, include_columns=True)
expected_full_output = "[{'user_id': 17, 'user_name': 'hwchase', 'user_bio': None}]" expected_full_output = "[{'user_id': 17, 'user_name': 'hwchase', 'user_bio': None}]"
assert full_output == expected_full_output assert full_output == expected_full_output
def test_sql_database_run_update() -> None: def test_sql_database_run_update(db: SQLDatabase) -> None:
"""Test commands which return no rows return an empty string.""" """Test commands which return no rows return an empty string."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine) # Provision.
stmt = insert(user).values(user_id=13, user_name="Harrison") stmt = insert(user).values(user_id=13, user_name="Harrison")
with engine.begin() as conn: db._execute(stmt)
conn.execute(stmt)
db = SQLDatabase(engine) # Query and verify.
command = "update user set user_name='Updated' where user_id = 13" command = "update user set user_name='Updated' where user_id = 13"
output = db.run(command) output = db.run(command)
expected_output = "" expected_output = ""
@ -198,7 +201,7 @@ def test_sql_database_run_update() -> None:
def test_sql_database_schema_translate_map() -> None: def test_sql_database_schema_translate_map() -> None:
"""Verify using statement-specific execution options.""" """Verify using statement-specific execution options."""
engine = create_engine("sqlite:///:memory:") engine = sa.create_engine("sqlite:///:memory:")
db = SQLDatabase(engine) db = SQLDatabase(engine)
# Define query using SQLAlchemy selectable. # Define query using SQLAlchemy selectable.

Loading…
Cancel
Save