Skip to content

Commit 199eb9a

Browse files
authored
Fix integration tests (#501)
## Summary This PR fixes the integration tests so they all pass. Previously, integration testing was an optional manual process. These tests would not always be in a working state and weren't always updated with associated source code changes. As the project matures, we need to rely on these tests and invest more in them. As a follow-up to this PR, I plan to run these tests as part of our PR pipeline to ensure they pass for all changes. This may make it more difficult for non-dbt Labs employees to contribute, but this is important to ensure the dbt MCP is stable and reliable. ## Checklist - [x] I have performed a self-review of my code - [x] I have made corresponding changes to the documentation (in https://github.com/dbt-labs/docs.getdbt.com) if required -- Mention it here - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes
1 parent bdcd886 commit 199eb9a

10 files changed

Lines changed: 211 additions & 207 deletions

File tree

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
kind: Under the Hood
2+
body: Fix integration tests
3+
time: 2026-01-06T18:00:13.319613-06:00

CONTRIBUTING.md

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,19 @@ Refer to `./cursor/rules` for standards and stylistic guidelines.
3030

3131
## Testing
3232

33-
This repo has automated tests which can be run with `task test:unit`. Additionally, there is a simple CLI tool which can be used to test by running `task client`. If you would like to test in a client like Cursor or Claude, use a configuration file like this:
33+
### Unit Testing
34+
35+
This repo has automated tests which can be run with `task test:unit`.
36+
37+
### Integration Testing
38+
39+
The integration tests exercise system interactions greater than just the unit tests. They require a dbt Platform environment that is setup with semantic layer, developer license, and PAT. These tests can be run with `task test:integration`.
40+
41+
For dbt Labs employees, a staging environment has been set up. Credentials for this environment can be found by searching for `dbt MCP Integration Test Credentials` in 1Password.
42+
43+
### Manual Testing
44+
45+
To test in a client like Cursor or Claude, use a configuration file like this:
3446

3547
```
3648
{

src/dbt_mcp/proxy/tools.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,7 @@ async def close(cls) -> None:
125125
async def register_proxied_tools(
126126
dbt_mcp: FastMCP,
127127
config_provider: ConfigProvider[ProxiedToolConfig],
128+
*,
128129
disabled_tools: set[ToolName],
129130
enabled_tools: set[ToolName],
130131
enabled_toolsets: set[Toolset],

src/remote_mcp/session.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,21 @@
88

99
@contextlib.asynccontextmanager
1010
async def session_context() -> AsyncGenerator[ClientSession, None]:
11+
host = os.environ.get("DBT_HOST")
12+
prefix = os.environ.get("MULTICELL_ACCOUNT_PREFIX")
13+
url = (
14+
f"https://{prefix}.{host}/api/ai/v1/mcp/"
15+
if prefix
16+
else f"https://{host}/api/ai/v1/mcp/"
17+
)
18+
token = os.environ.get("DBT_TOKEN")
19+
prod_environment_id = os.environ.get("DBT_PROD_ENV_ID", "")
1120
async with (
1221
streamablehttp_client(
13-
url=f"https://{os.environ.get('DBT_HOST')}/api/ai/v1/mcp/",
22+
url=url,
1423
headers={
15-
"Authorization": f"token {os.environ.get('DBT_TOKEN')}",
16-
"x-dbt-prod-environment-id": os.environ.get("DBT_PROD_ENV_ID", ""),
24+
"Authorization": f"token {token}",
25+
"x-dbt-prod-environment-id": prod_environment_id,
1726
},
1827
) as (
1928
read_stream,
Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
import os
2+
3+
import pytest
4+
5+
from dbt_mcp.config.config_providers import (
6+
ConfigProvider,
7+
DefaultDiscoveryConfigProvider,
8+
DiscoveryConfig,
9+
)
10+
from dbt_mcp.config.settings import CredentialsProvider, DbtMcpSettings
11+
from dbt_mcp.discovery.client import (
12+
DEFAULT_MAX_NODE_QUERY_LIMIT,
13+
DEFAULT_PAGE_SIZE,
14+
ExposuresFetcher,
15+
MetadataAPIClient,
16+
ModelsFetcher,
17+
PaginatedResourceFetcher,
18+
SourcesFetcher,
19+
)
20+
21+
22+
@pytest.fixture
23+
def config_provider() -> ConfigProvider[DiscoveryConfig]:
24+
# Set up environment variables needed by DbtMcpSettings
25+
host = os.getenv("DBT_HOST")
26+
token = os.getenv("DBT_TOKEN")
27+
prod_env_id = os.getenv("DBT_PROD_ENV_ID")
28+
29+
if not host or not token or not prod_env_id:
30+
raise ValueError(
31+
"DBT_HOST, DBT_TOKEN, and DBT_PROD_ENV_ID environment variables are "
32+
"required"
33+
)
34+
35+
# DbtMcpSettings will automatically pick up from environment variables
36+
settings = DbtMcpSettings() # type: ignore
37+
credentials_provider = CredentialsProvider(settings)
38+
return DefaultDiscoveryConfigProvider(credentials_provider)
39+
40+
41+
@pytest.fixture
42+
def api_client(config_provider: ConfigProvider[DiscoveryConfig]) -> MetadataAPIClient:
43+
return MetadataAPIClient(config_provider)
44+
45+
46+
@pytest.fixture
47+
def models_fetcher(api_client: MetadataAPIClient) -> ModelsFetcher:
48+
paginator = PaginatedResourceFetcher(
49+
api_client,
50+
edges_path=("data", "environment", "applied", "models", "edges"),
51+
page_info_path=("data", "environment", "applied", "models", "pageInfo"),
52+
page_size=DEFAULT_PAGE_SIZE,
53+
max_node_query_limit=DEFAULT_MAX_NODE_QUERY_LIMIT,
54+
)
55+
return ModelsFetcher(api_client, paginator=paginator)
56+
57+
58+
@pytest.fixture
59+
def exposures_fetcher(api_client: MetadataAPIClient) -> ExposuresFetcher:
60+
paginator = PaginatedResourceFetcher(
61+
api_client,
62+
edges_path=("data", "environment", "definition", "exposures", "edges"),
63+
page_info_path=(
64+
"data",
65+
"environment",
66+
"definition",
67+
"exposures",
68+
"pageInfo",
69+
),
70+
page_size=DEFAULT_PAGE_SIZE,
71+
max_node_query_limit=DEFAULT_MAX_NODE_QUERY_LIMIT,
72+
)
73+
return ExposuresFetcher(api_client, paginator=paginator)
74+
75+
76+
@pytest.fixture
77+
def sources_fetcher(api_client: MetadataAPIClient) -> SourcesFetcher:
78+
paginator = PaginatedResourceFetcher(
79+
api_client,
80+
edges_path=("data", "environment", "applied", "sources", "edges"),
81+
page_info_path=("data", "environment", "applied", "sources", "pageInfo"),
82+
page_size=DEFAULT_PAGE_SIZE,
83+
max_node_query_limit=DEFAULT_MAX_NODE_QUERY_LIMIT,
84+
)
85+
return SourcesFetcher(api_client, paginator=paginator)

tests/integration/discovery/test_discovery.py

Lines changed: 9 additions & 110 deletions
Original file line numberDiff line numberDiff line change
@@ -1,83 +1,15 @@
1-
import os
2-
31
import pytest
42

5-
from dbt_mcp.config.config_providers import DefaultDiscoveryConfigProvider
6-
from dbt_mcp.config.settings import CredentialsProvider, DbtMcpSettings
3+
from dbt_mcp.config.config_providers import ConfigProvider, DiscoveryConfig
74
from dbt_mcp.discovery.client import (
8-
DEFAULT_MAX_NODE_QUERY_LIMIT,
95
DEFAULT_PAGE_SIZE,
106
ExposuresFetcher,
11-
MetadataAPIClient,
127
ModelFilter,
138
ModelsFetcher,
14-
PaginatedResourceFetcher,
159
SourcesFetcher,
1610
)
1711
from dbt_mcp.discovery.tools import DISCOVERY_TOOLS, DiscoveryToolContext
18-
19-
20-
@pytest.fixture
21-
def api_client() -> MetadataAPIClient:
22-
# Set up environment variables needed by DbtMcpSettings
23-
host = os.getenv("DBT_HOST")
24-
token = os.getenv("DBT_TOKEN")
25-
prod_env_id = os.getenv("DBT_PROD_ENV_ID")
26-
27-
if not host or not token or not prod_env_id:
28-
raise ValueError(
29-
"DBT_HOST, DBT_TOKEN, and DBT_PROD_ENV_ID environment variables are required"
30-
)
31-
32-
# Create settings and credentials provider
33-
# DbtMcpSettings will automatically pick up from environment variables
34-
settings = DbtMcpSettings() # type: ignore
35-
credentials_provider = CredentialsProvider(settings)
36-
config_provider = DefaultDiscoveryConfigProvider(credentials_provider)
37-
38-
return MetadataAPIClient(config_provider)
39-
40-
41-
@pytest.fixture
42-
def models_fetcher(api_client: MetadataAPIClient) -> ModelsFetcher:
43-
paginator = PaginatedResourceFetcher(
44-
api_client,
45-
edges_path=("data", "environment", "applied", "models", "edges"),
46-
page_info_path=("data", "environment", "applied", "models", "pageInfo"),
47-
page_size=DEFAULT_PAGE_SIZE,
48-
max_node_query_limit=DEFAULT_MAX_NODE_QUERY_LIMIT,
49-
)
50-
return ModelsFetcher(api_client, paginator=paginator)
51-
52-
53-
@pytest.fixture
54-
def exposures_fetcher(api_client: MetadataAPIClient) -> ExposuresFetcher:
55-
paginator = PaginatedResourceFetcher(
56-
api_client,
57-
edges_path=("data", "environment", "definition", "exposures", "edges"),
58-
page_info_path=(
59-
"data",
60-
"environment",
61-
"definition",
62-
"exposures",
63-
"pageInfo",
64-
),
65-
page_size=DEFAULT_PAGE_SIZE,
66-
max_node_query_limit=DEFAULT_MAX_NODE_QUERY_LIMIT,
67-
)
68-
return ExposuresFetcher(api_client, paginator=paginator)
69-
70-
71-
@pytest.fixture
72-
def sources_fetcher(api_client: MetadataAPIClient) -> SourcesFetcher:
73-
paginator = PaginatedResourceFetcher(
74-
api_client,
75-
edges_path=("data", "environment", "applied", "sources", "edges"),
76-
page_info_path=("data", "environment", "applied", "sources", "pageInfo"),
77-
page_size=DEFAULT_PAGE_SIZE,
78-
max_node_query_limit=DEFAULT_MAX_NODE_QUERY_LIMIT,
79-
)
80-
return SourcesFetcher(api_client, paginator=paginator)
12+
from dbt_mcp.tools.tool_names import ToolName
8113

8214

8315
@pytest.mark.asyncio
@@ -91,17 +23,10 @@ async def test_fetch_models(models_fetcher: ModelsFetcher):
9123
# Validate structure of returned models
9224
for model in results:
9325
assert "name" in model
94-
assert "compiledCode" in model
26+
assert "uniqueId" in model
27+
assert "description" in model
9528
assert isinstance(model["name"], str)
9629

97-
# If catalog exists, validate its structure
98-
if model.get("catalog"):
99-
assert isinstance(model["catalog"], dict)
100-
if "columns" in model["catalog"]:
101-
for column in model["catalog"]["columns"]:
102-
assert "name" in column
103-
assert "type" in column
104-
10530

10631
@pytest.mark.asyncio
10732
async def test_fetch_models_with_filter(models_fetcher: ModelsFetcher):
@@ -209,7 +134,7 @@ async def test_fetch_exposures_pagination(exposures_fetcher: ExposuresFetcher):
209134
assert isinstance(results, list)
210135

211136
# If we have more than the page size, ensure no duplicates
212-
if len(results) > 100: # PAGE_SIZE is 100
137+
if len(results) > DEFAULT_PAGE_SIZE:
213138
unique_ids = set()
214139
for exposure in results:
215140
unique_id = exposure["uniqueId"]
@@ -246,17 +171,6 @@ async def test_fetch_sources(sources_fetcher: SourcesFetcher):
246171
if source.get("freshness"):
247172
freshness = source["freshness"]
248173
assert isinstance(freshness, dict)
249-
# These fields may be present depending on configuration
250-
if "freshnessStatus" in freshness:
251-
assert isinstance(freshness["freshnessStatus"], str)
252-
if "maxLoadedAt" in freshness:
253-
assert freshness["maxLoadedAt"] is None or isinstance(
254-
freshness["maxLoadedAt"], str
255-
)
256-
if "maxLoadedAtTimeAgoInS" in freshness:
257-
assert freshness["maxLoadedAtTimeAgoInS"] is None or isinstance(
258-
freshness["maxLoadedAtTimeAgoInS"], int
259-
)
260174

261175

262176
@pytest.mark.asyncio
@@ -283,33 +197,18 @@ async def test_fetch_sources_with_filter(sources_fetcher: SourcesFetcher):
283197

284198

285199
@pytest.mark.asyncio
286-
async def test_get_all_sources_tool():
200+
async def test_get_all_sources_tool(
201+
config_provider: ConfigProvider[DiscoveryConfig],
202+
) -> None:
287203
"""Test the get_all_sources tool function integration."""
288-
from dbt_mcp.config.config_providers import DefaultDiscoveryConfigProvider
289-
from dbt_mcp.config.settings import CredentialsProvider, DbtMcpSettings
290-
291-
# Set up environment variables needed by DbtMcpSettings
292-
host = os.getenv("DBT_HOST")
293-
token = os.getenv("DBT_TOKEN")
294-
prod_env_id = os.getenv("DBT_PROD_ENV_ID")
295-
296-
if not host or not token or not prod_env_id:
297-
pytest.skip(
298-
"DBT_HOST, DBT_TOKEN, and DBT_PROD_ENV_ID environment variables are required"
299-
)
300-
301-
# Create settings and config provider
302-
settings = DbtMcpSettings() # type: ignore
303-
credentials_provider = CredentialsProvider(settings)
304-
config_provider = DefaultDiscoveryConfigProvider(credentials_provider)
305204

306205
# Create tool definitions
307206
tool_definitions = DISCOVERY_TOOLS
308207

309208
# Find the get_all_sources tool
310209
get_all_sources_tool = None
311210
for tool_def in tool_definitions:
312-
if tool_def.get_name() == "get_all_sources":
211+
if tool_def.get_name() == ToolName.GET_ALL_SOURCES:
313212
get_all_sources_tool = tool_def
314213
break
315214

0 commit comments

Comments
 (0)