Add LangChain Lakebase integration tests by dhruv0811 · Pull Request #346 · databricks/databricks-ai-bridge

dhruv0811 · 2026-02-24T19:43:39Z

Summary

Added LangChain Lakebase integration tests covering DatabricksStore (sync + async) and CheckpointSaver (sync + async) wrappers. Tests validate the bridge code's behavior — pool creation, connection borrowing, _with_store handoff, and async with lifecycle management — against a live Databricks Lakebase instance.

Test plan

All 12 tests collected and skip gracefully without LAKEBASE_INSTANCE_NAME
Existing unit tests (test_store.py, test_checkpoint.py) still pass (16/16)
Core Lakebase tests pass (54/54)
LangChain Lakebase integration tests pass against live Lakebase instance via ai-oss runner: https://github.com/databricks-eng/ai-oss-integration-tests-runner/actions/runs/22643496765

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…ckpoint_ns - Make cleanup fixtures module-scoped to avoid dropping tables between tests (PostgresStore/Saver track migrations internally) - Add missing checkpoint_ns to all checkpoint config dicts - Remove stale table-tracking append/extend calls Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

- Drop store_migrations, vector_migrations, store_vectors tables in cleanup fixture (PostgresStore skips setup if migration table exists) - Add auth_type="pat" to create_workspace_client_with_token to avoid oauth+pat conflict when DATABRICKS_CLIENT_ID env var is set Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…igration state

tests/integration_tests/lakebase/test_lakebase_integration.py

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py

jennsun

these tests are looking great! my main request is to add a test that tests the core value databricksstore brings to agents (semantic search)

https://github.com/databricks/databricks-ai-bridge/pull/346/changes#r2850132619

…ch test, ordering assertions

…e list order)

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py

jennsun

thanks for the updates! can we make sure we have coverage/parallel tests for both sync + async versions of the lakebase langchain classes? good to merge after those are synced :)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…__/__exit__) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py

…pool closed Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

… in put_and_get Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py

…helper, fixture docs - Move all imports to top level with pytest.importorskip("langgraph") guard - Extract _make_checkpoint() helper shared across sync/async tests - Clarify scope="module" behavior in fixture docstrings Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

dhruv0811 and others added 7 commits February 24, 2026 11:43

Add LangChain Lakebase integration tests (Store + Checkpoint)

2919991

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Fix ruff formatting: add blank line before class definition

0b24b3d

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Drop store/checkpoint tables before AND after tests to handle stale m…

63a105f

…igration state

Add async search and delete store tests for parity with sync

1d002fe

Add async checkpoint list test for parity with sync

6073b70

dhruv0811 requested review from annzhang-db and jennsun February 24, 2026 21:26

jennsun reviewed Feb 24, 2026

View reviewed changes

tests/integration_tests/lakebase/test_lakebase_integration.py Outdated Show resolved Hide resolved

jennsun reviewed Feb 24, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Outdated Show resolved Hide resolved

jennsun reviewed Feb 25, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Show resolved Hide resolved

jennsun reviewed Feb 25, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Outdated Show resolved Hide resolved

jennsun reviewed Feb 25, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Outdated Show resolved Hide resolved

jennsun reviewed Feb 25, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Show resolved Hide resolved

jennsun requested changes Feb 25, 2026

View reviewed changes

dhruv0811 added 3 commits February 25, 2026 15:30

Address PR review: use OAuth M2M, LakebaseClient cleanup, vector sear…

d97ef83

…ch test, ordering assertions

Remove checkpoint list ordering assertion (ts field does not determin…

1d77704

…e list order)

Fix ruff formatting: single-line function signature

7d023c8

dhruv0811 requested a review from jennsun February 25, 2026 23:47

jennsun reviewed Feb 26, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Show resolved Hide resolved

jennsun reviewed Feb 26, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Outdated Show resolved Hide resolved

jennsun reviewed Feb 26, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Show resolved Hide resolved

jennsun reviewed Feb 26, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Show resolved Hide resolved

jennsun approved these changes Feb 26, 2026

View reviewed changes

dhruv0811 and others added 2 commits February 26, 2026 11:37

Sync async test parity

15981f5

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Remove sync store context manager test (DatabricksStore lacks __enter…

c71bd34

…__/__exit__) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

annzhang-db reviewed Mar 3, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Outdated Show resolved Hide resolved

dhruv0811 and others added 2 commits March 3, 2026 11:39

Combine checkpoint context manager tests with write_and_read, assert …

42554ff

…pool closed Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Remove redundant async store context manager test, assert pool closed…

ef3627f

… in put_and_get Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

dhruv0811 requested a review from annzhang-db March 3, 2026 19:45

bbqiu reviewed Mar 3, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Show resolved Hide resolved

bbqiu reviewed Mar 3, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Outdated Show resolved Hide resolved

bbqiu reviewed Mar 3, 2026

View reviewed changes

integrations/langchain/tests/integration_tests/test_langchain_lakebase.py Show resolved Hide resolved

dhruv0811 and others added 2 commits March 3, 2026 13:27

Fix ruff import sorting

d5d5b18

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add LangChain Lakebase integration tests#346

Add LangChain Lakebase integration tests#346
dhruv0811 wants to merge 16 commits intomainfrom
lakebase-langchain-integration-tests

dhruv0811 commented Feb 24, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jennsun left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jennsun left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

dhruv0811 commented Feb 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jennsun left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jennsun left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

dhruv0811 commented Feb 24, 2026 •

edited

Loading