fix: allow _with_context tests to collect sample rows

[joostboon]

Summary

• _with_context tests are Elementary-namespaced but behave like regular dbt tests (they return failing rows rather than handling their own result row collection)
• The test materialization was early-returning for all elementary-namespaced tests, bypassing the sampling logic entirely
• This meant _with_context tests never wrote rows to test_result_rows, even when tagged with show_sample_rows

Fix

Check if the test name ends with _with_context before early-returning, allowing these tests to go through the standard handle_dbt_test sampling path.

Test plan

• Run a failing _with_context test (e.g. not_null_with_context) and verify rows appear in test_result_rows
• Verify non-_with_context Elementary tests (anomaly, schema change) still early-return and handle their own result rows

🤖 Generated with Claude Code

Summary by CodeRabbit

• Bug Fixes
  • Improved handling of elementary tests by applying specialized processing only to recognized test types.
  • Ensured other elementary-namespaced tests, including context-aware tests, follow the standard processing flow.
  • Fixed sampling for not_null_with_context tests so configured sample counts are respected and context values are preserved in sampled results.

[github-actions]

👋 @joostboon
Thank you for raising your pull request.
Please make sure to add tests and document all user-facing changes.
You can do this by editing the docs files in the elementary repository.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

The materialize_test hook now returns early only for recognized elementary test types. An integration test verifies sampled not_null_with_context results, including failure status, sample count, and context values.

Changes

Test materialization behavior

Layer / File(s)	Summary
Classify elementary tests macros/edr/materializations/test/test.sql	The hook lowercases the test name, resolves the elementary test type, and returns the base materialization only for recognized elementary tests. Other elementary-namespaced tests continue through the normal flow.
Validate sampled context results integration_tests/tests/test_with_context_sampling.py	The integration test enables sampling, checks the failure state, loads sampled rows, and verifies the sample count and context values.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly describes the main change: enabling Elementary _with_context tests to collect sample rows.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/with-context-tests-sample-rows

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

This pull request is stale because it has been open for too long with no activity.
If you would like the pull request to remain open, please remove the stale label or leave a comment.

[github-actions]

This pull request was closed because it has been inactive for too long while being marked as stale.
If you would like the pull request to reopen, please leave a comment.

[joostboon]

Reopening — this was closed by stale-bot on 2026-07-10 after no human review, not on its merits. Added the integration test the contribution bot asked for, and independently re-verified the bug against a live Snowflake warehouse.

The bug is in every release that has the feature

The namespace early-return landed in 0.23.0; the _with_context tests landed in 0.24.0. So they have never collected sample rows in any released version.

version	_with_context tests	namespace early-return
0.22.0	absent	no
0.23.0	absent	yes
0.24.0 → 0.25.1	present	yes

Verified end to end

elementary 0.25.1, dbt-snowflake 1.11.1, using elementary.not_null_with_context exactly as documented on a column with 55 nulls and five context_columns:

	sample rows in test_result_rows	failed_row_count
stock 0.25.1	0	null
with this branch	5	55

Note failed_row_count was also null before, since calculate_failed_count sits after the early return — so these tests were skipping that too.

A useful corroborating detail: the package already classifies these as regular dbt tests. get_elementary_test_type returns a type only for the anomaly and schema-change families, so get_test_type returns dbt_test for _with_context — which is what lands in elementary_test_results.test_type. The materialization branching on namespace instead of test type was inconsistent with the package's own classification, which is what this fixes.

What I added

integration_tests/tests/test_with_context_sampling.py — asserts an elementary-namespaced _with_context test writes test_sample_row_count rows and that the requested context column travels with each failing row. There were no integration tests covering the _with_context family at all, which is why this shipped unnoticed. Modelled on test_sampling.py and uses the same harness calls. It fails against the unfixed materialization.

I have not run the integration suite myself — no warehouse credentials for the harness — so CI is the first real execution of that test.

🤖 Generated with Claude Code

[coderabbitai]

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@integration_tests/tests/test_with_context_sampling.py`:
- Around line 44-46: Update the assertions in the sample-validation loop to
compare each sample’s CONTEXT_COLUMN value against the exact expected context
value derived from null_count, rather than only checking the "context-" prefix;
keep the existing null assertion for COLUMN_NAME.
- Line 36: Update the test assertion near the existing status check to validate
that test_result["failed_row_count"] equals null_count. Keep the existing
failure-status assertion and use the input’s null_count value as the expected
failed-row count.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 8b82eec2-717e-4ec8-981e-412932f557b1

📥 Commits

Reviewing files that changed from the base of the PR and between 37988b0 and 5ffffc4.

📒 Files selected for processing (1)

• integration_tests/tests/test_with_context_sampling.py

[joostboon]

CI is red, but none of it is this PR. Breakdown:

job	cause	related?
latest_official, snowflake	Database Error 251006: Password is empty — failed at dbt debug, before any test ran	no
fusion, snowflake	[Snowflake] 260002: password is empty	no
latest_official, athena	InvalidClientTokenId / UnrecognizedClientException — expired AWS credentials	no
latest_official, sqlserver	errors in test_exposure_schema_validity.py	no
fusion, bigquery	FAILED test_exposure_schema_validity_no_exposures	no

Two independent pre-existing issues: stale CI credentials for Snowflake and Athena (expected on a branch idle since April), and test_exposure_schema_validity.py failing on sqlserver and bigquery-fusion.

The new test passed on every adapter that could connect:

postgres  [gw2] [ 98%] PASSED test_with_context_sampling.py::test_with_context_test_is_sampled
duckdb    [gw6] [ 98%] PASSED test_with_context_sampling.py::test_with_context_test_is_sampled
redshift  [gw0] [ 98%] PASSED test_with_context_sampling.py::test_with_context_test_is_sampled

plus the other green adapters (fabric, dremio, vertica, clickhouse, databricks_catalog, latest_pre/postgres).

Snowflake CI couldn't run, but that is the adapter I verified the fix on by hand — elementary 0.25.1 / dbt-snowflake 1.11.1, elementary.not_null_with_context on a column with 55 nulls: 0 sample rows before, 5 after, with every requested context column attached, and failed_row_count going from null to 55.

Someone with access will need to refresh the Snowflake and Athena CI secrets to get a fully green run.

🤖 Generated with Claude Code

[joostboon]

Both CodeRabbit suggestions applied in 4b176a8 — both were right, and the first is more load-bearing than it looks.

failed_row_count — applied. Worth noting why it matters here rather than being a generic nicety: calculate_failed_count sits after the same early return the sampling did, so it was skipped too. Measured on Snowflake with 0.25.1, elementary.not_null_with_context on a column with 55 nulls:

	sample rows	failed_row_count
unfixed	0	null
fixed	5	55

So this assertion covers a second broken path, not just the samples. Consistent with test_failed_row_count.py, which asserts the same key the same way.

Exact context values — applied. Used set membership rather than equality, since sampling returns an arbitrary test_sample_row_count of the null_count failing rows, so no single expected value would hold. This does prove the materialization preserved the source value, which the startswith check did not.

Also merged current master into the branch (68 commits behind). No conflicts, and the gate survived intact — net diff against master is still just test.sql (+10/-2) and the new test. I verified the harness API this test depends on (dbt_project.test, samples_query, run_query) is unchanged across those commits.

🤖 Generated with Claude Code

[haritamar]

Can we perhaps instead be more explicit, and return a type in get_elementary_test_type for the relevant tests, and then check here that it's not that type?

[joostboon]

Done in c33465f — agreed, explicit is better. get_elementary_test_type now returns a "with_context" type for the six tests in the family, and the materialization checks elementary_test_type != "with_context" instead of leaning on the macro returning none. It also means a future elementary test family defaults to the existing skip behaviour rather than silently landing in this bug again.

One thing worth knowing, since it caught me while implementing it: the new type must not reach test_type. All three alerts models partition on that value exactly —

• alerts_dbt_tests.sql:41 → test_type = 'dbt_test'
• alerts_anomaly_detection.sql:41 → test_type = 'anomaly_detection'
• alerts_schema_changes.sql:47 → test_type = 'schema_change'

so a fourth value would match none of them and with_context test failures would raise no alerts at all. get_test_type therefore maps with_context → dbt_test, keeping the reported type and the alerts behaviour exactly as before; the new type stays internal to the materialization gate. I've added assert test_result["test_type"] == "dbt_test" to the integration test so that can't regress silently.

Happy to change the type name, or to move the exemption list somewhere more visible than get_elementary_test_type, if you'd prefer it shaped differently.

🤖 Generated with Claude Code