Test pruner and refactor newDB

[mkysel]

Refactor database cleanup handling in tests to use testing framework's t.Cleanup() mechanism instead of explicit cleanup functions

• Removes explicit cleanup function calls across multiple test files, replacing them with the testing framework's built-in t.Cleanup() mechanism in store.go
• Modifies NewDB and NewDBs functions in store.go to no longer return cleanup functions
• Adds comprehensive test suite for pruning functionality in prune_test.go, covering expired envelopes, dry run mode, concurrent locks, and large payload scenarios

📍Where to Start

Start with the core changes to database cleanup handling in store.go, particularly the openDB function which now registers cleanup with t.Cleanup() instead of returning a cleanup function.

---

Macroscope summarized d1a4b2a.

Summary by CodeRabbit

• Tests
  • Updated test setup and teardown procedures to simplify resource management by relying on automatic cleanup, removing explicit cleanup function calls across multiple test files.
  • Refactored helper and setup functions in tests to streamline usage and reduce manual cleanup handling.
  • Added comprehensive new unit tests for database pruning functionality, including scenarios for concurrency, dry-run mode, and large datasets.

[graphite-app]

How to use the Graphite Merge Queue

Add either label to this PR to merge it via the merge queue:

• Queue - adds this PR to the back of the merge queue
• Hotfix - for urgent hot fixes, skip the queue and merge this PR next

You must have a Graphite account in order to use the merge queue. Sign up using this link.

An organization admin has enabled the Graphite Merge Queue in this repository.

Please do not merge from GitHub as this will restart CI on PRs being processed by the merge queue.

[coderabbitai]

Walkthrough

The changes remove explicit cleanup function returns and calls from test database setup utilities and all dependent test files, instead relying on Go's t.Cleanup mechanism for resource management. Function signatures and usages are updated accordingly across many test files. Additionally, a new test suite for the prune package is introduced, providing comprehensive tests for envelope pruning logic.

Changes

File(s)	Change Summary
pkg/testutils/store.go	Refactored database setup functions to register cleanup with t.Cleanup instead of returning explicit cleanup functions; updated all signatures to remove cleanup returns.
pkg/blockchain/nonceManager_test.go pkg/db/gatewayEnvelope_test.go pkg/db/queries_test.go pkg/db/sequences_test.go pkg/db/unsettledUsage_test.go pkg/fees/calculator_test.go pkg/indexer/blockTracker_test.go pkg/payerreport/manager_test.go	Removed usage and deferral of cleanup functions returned by testutils.NewDB; tests now only capture the DB handle. No other logic changed.
pkg/db/congestion_test.go	Updated setupTest to remove cleanup return; all call sites updated to omit cleanup handling.
pkg/db/subscription_test.go	Updated setup to remove cleanup return; all call sites updated to omit cleanup handling.
pkg/indexer/e2e_test.go pkg/indexer/storer/groupMessage_test.go pkg/indexer/storer/identityUpdate_test.go	Removed capturing and invocation of cleanup functions from test DB setup; updated helper/teardown logic accordingly.
pkg/registrant/registrant_test.go	Refactored test setup helpers to no longer return cleanup functions; updated all test usages to omit cleanup handling.
pkg/server/server_test.go	Removed deferred cleanup calls from DB setup; only DB handles are used.
pkg/testutils/api/api.go	Modified NewTestAPIServer to not capture or invoke DB cleanup; cleanup function no longer calls DB cleanup.
pkg/prune/prune_test.go	New file: Introduces comprehensive unit tests for the prune executor, including helpers for test data setup, lock simulation, and validation of pruning logic under various scenarios.

Sequence Diagram(s)

sequenceDiagram
    participant TestFunction
    participant testutils.NewDB
    participant t.Cleanup
    TestFunction->>testutils.NewDB: Create test DB
    testutils.NewDB->>t.Cleanup: Register DB cleanup callback
    TestFunction-->>testutils.NewDB: Receives DB handle only
    Note right of TestFunction: No explicit cleanup function returned or deferred

Loading

sequenceDiagram
    participant Test as Prune Test
    participant setupTestData
    participant makeTestExecutor
    participant prune.Executor
    participant DB

    Test->>setupTestData: Insert expired & valid envelopes
    Test->>makeTestExecutor: Create prune executor
    Test->>prune.Executor: Run()
    prune.Executor->>DB: Delete expired envelopes
    Test->>DB: Query remaining envelopes
    Test->>Test: Assert expected state

Loading

Suggested reviewers

• fbac

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (1.64.8)

Error: you are using a configuration file for golangci-lint v2 with golangci-lint v1: please use golangci-lint v2
Failed executing command with error: you are using a configuration file for golangci-lint v2 with golangci-lint v1: please use golangci-lint v2

Tip

⚡️ Faster reviews with caching

• CodeRabbit now supports caching for code and dependencies, helping speed up reviews. This means quicker feedback, reduced wait times, and a smoother review experience overall. Cached data is encrypted and stored securely. This feature will be automatically enabled for all accounts on May 16th. To opt out, configure Review - Disable Cache at either the organization or repository level. If you prefer to disable all data retention across your organization, simply turn off the Data Retention setting under your Organization Settings.

Enjoy the performance boost—your workflow just got faster.

---

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 80a852a and d1a4b2a.

📒 Files selected for processing (18)

• pkg/blockchain/nonceManager_test.go (4 hunks)
• pkg/db/congestion_test.go (4 hunks)
• pkg/db/gatewayEnvelope_test.go (4 hunks)
• pkg/db/queries_test.go (3 hunks)
• pkg/db/sequences_test.go (13 hunks)
• pkg/db/subscription_test.go (4 hunks)
• pkg/db/unsettledUsage_test.go (2 hunks)
• pkg/fees/calculator_test.go (1 hunks)
• pkg/indexer/blockTracker_test.go (4 hunks)
• pkg/indexer/e2e_test.go (1 hunks)
• pkg/indexer/storer/groupMessage_test.go (1 hunks)
• pkg/indexer/storer/identityUpdate_test.go (1 hunks)
• pkg/payerreport/manager_test.go (1 hunks)
• pkg/prune/prune_test.go (1 hunks)
• pkg/registrant/registrant_test.go (11 hunks)
• pkg/server/server_test.go (2 hunks)
• pkg/testutils/api/api.go (1 hunks)
• pkg/testutils/store.go (1 hunks)

🧰 Additional context used

🧬 Code Graph Analysis (17)

pkg/fees/calculator_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/db/sequences_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/payerreport/manager_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/db/unsettledUsage_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/indexer/storer/identityUpdate_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/db/gatewayEnvelope_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/indexer/storer/groupMessage_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/blockchain/nonceManager_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/indexer/e2e_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/server/server_test.go (1)

pkg/testutils/store.go (1)

• NewDBs (72-82)

pkg/db/subscription_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/db/congestion_test.go (2)

pkg/db/queries/db.go (2)

• Queries (23-25)
• New (19-21)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/testutils/api/api.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/indexer/blockTracker_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/db/queries_test.go (1)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/registrant/registrant_test.go (4)

pkg/testutils/log.go (1)

• NewLog (17-25)

pkg/mocks/registry/mock_NodeRegistry.go (1)

• NewMockNodeRegistry (291-301)

pkg/testutils/store.go (1)

• NewDB (65-70)

pkg/registrant/registrant.go (1)

• Registrant (26-30)

pkg/prune/prune_test.go (6)

pkg/db/queries/db.go (1)

• New (19-21)

pkg/db/queries/envelopes.sql.go (1)

• InsertGatewayEnvelopeParams (97-105)

pkg/config/pruneOptions.go (1)

• PruneConfig (3-6)

pkg/prune/prune.go (2)

• Executor (16-21)
• NewPruneExecutor (23-35)

pkg/testutils/log.go (1)

• NewLog (17-25)

pkg/testutils/store.go (1)

• NewDBs (72-82)

⏰ Context from checks skipped due to timeout of 90000ms (5)

• GitHub Check: Code Review
• GitHub Check: Push Docker Images to GitHub Packages (xmtpd-cli)
• GitHub Check: Push Docker Images to GitHub Packages (xmtpd)
• GitHub Check: Test (Node)
• GitHub Check: Upgrade Tests

🔇 Additional comments (43)

pkg/prune/prune_test.go (9)

18-21: Good use of constants for test configuration values.

The constants DEFAULT_EXPIRED_CNT and DEFAULT_VALID_CNT help make the tests more maintainable by centralizing these values.

---

23-60: Well-structured test data setup function.

The setupTestData function creates both expired and valid envelopes with appropriate test data. It properly handles error checking and creates a realistic test environment.

---

62-74: Good helper function for test executor creation.

The makeTestExecutor function provides a clean way to create test executor instances with different configurations, promoting code reuse across test cases.

---

76-101: Comprehensive test for basic pruning functionality.

This test effectively verifies that:

1. Expired envelopes are properly pruned
2. Non-expired envelopes remain untouched
3. The pruning operation completes without errors

The assertions are clear and include descriptive error messages.

---

103-132: Thorough verification of dry-run behavior.

This test correctly verifies that dry-run mode doesn't actually delete any envelopes while still completing successfully. The assertions check both expired and non-expired envelope counts.

---

134-146: Effective helper function for concurrency testing.

The openAndHoldLock function simulates a real-world scenario where a row might be locked by another transaction, which is important for testing how the pruner handles concurrent access.

---

148-180: Good test for handling row locks during pruning.

This test verifies that the pruner:

1. Works correctly when some rows are locked
2. Successfully prunes unlocked rows
3. Can finish pruning previously locked rows once they become available

This is important for ensuring the pruner is resilient in a production environment.

---

182-198: Useful helper function for verification.

The getRemainingSequenceIds function provides a clean way to verify which specific rows remain in the database after pruning operations, enabling more detailed assertions.

---

200-228: Effective test for large payload pruning with multiple cycles.

This test verifies that the pruner:

1. Respects the MaxCycles configuration
2. Successfully continues pruning across multiple runs
3. Eventually removes all expired envelopes when cycle limits are in place

This is important for ensuring the pruner can handle large volumes of data without overwhelming system resources.

pkg/db/sequences_test.go (2)

20-20: Simplification of test resource management.

The update removes explicit cleanup function captures, instead relying on Go's t.Cleanup mechanism now implemented in the testutils.NewDB function. This simplifies the test code while maintaining proper resource cleanup.

---

33-33: Consistent application of resource management simplification.

Similar changes are made throughout the file, consistently removing explicit cleanup function captures while maintaining the core test logic. This follows Go's best practices for test cleanup and makes the tests more readable.

Also applies to: 87-87, 139-139, 164-164, 189-189, 216-216, 252-252, 273-273, 294-294, 332-332, 375-375, 408-408

pkg/blockchain/nonceManager_test.go (1)

98-98: Simplified test resource management.

The changes consistently remove explicit cleanup function captures throughout the file, instead relying on Go's automatic cleanup mechanism now implemented in testutils.NewDB. This improves code readability while maintaining proper resource management.

Also applies to: 117-117, 137-137, 158-158

pkg/payerreport/manager_test.go (1)

21-21: Simplified resource management in setup function.

The update removes explicit cleanup function capture, relying on Go's t.Cleanup mechanism now implemented in the testutils.NewDB function. This makes the code cleaner while maintaining proper test cleanup.

pkg/testutils/api/api.go (1)

105-105: LGTM: Successfully refactored to use Go's t.Cleanup mechanism.

The code now correctly uses the new signature of testutils.NewDB that returns only a database connection and DSN, without an explicit cleanup function. This aligns with the refactoring to use Go's built-in t.Cleanup mechanism for resource management.

pkg/db/queries_test.go (1)

35-35: LGTM: Consistently removed explicit cleanup handling.

The test functions have been correctly updated to use the new signature of testutils.NewDB that returns only the database connection and DSN, without an explicit cleanup function. This change is consistent across all test functions in this file and aligns with the project-wide refactoring to use Go's t.Cleanup mechanism.

Also applies to: 89-89, 140-140

pkg/indexer/blockTracker_test.go (1)

21-21: LGTM: Tests properly refactored for cleanup handling.

All test functions have been consistently updated to use the new signature of testutils.NewDB, removing explicit cleanup function handling. This change is well-aligned with the project's transition to using Go's t.Cleanup mechanism for resource management.

Also applies to: 36-36, 77-77, 125-125

pkg/server/server_test.go (1)

76-76: LGTM: Properly adapted to use the new NewDBs function.

The test functions now correctly use the updated testutils.NewDBs function that returns only the database connections without an explicit cleanup function. This change is consistent with the project's move to use Go's t.Cleanup mechanism for resource management.

Also applies to: 219-219

pkg/fees/calculator_test.go (1)

74-74: Refactoring cleanup to use Go's t.Cleanup mechanism

The code now uses Go's built-in t.Cleanup mechanism instead of explicitly capturing and deferring a cleanup function. This is part of a codebase-wide refactoring that simplifies test resource management.

pkg/db/gatewayEnvelope_test.go (1)

42-42: Consistent cleanup handling refactoring across test files

These changes follow the same pattern of removing explicit cleanup function capture and calls. The testutils.NewDB function has been updated to handle cleanup internally using t.Cleanup, which simplifies test code by removing the need for explicit cleanup management.

Also applies to: 79-79, 98-98, 152-152

pkg/indexer/storer/identityUpdate_test.go (2)

28-28: Refactored database cleanup in test helper

The test helper no longer captures a cleanup function from testutils.NewDB, aligning with the codebase-wide refactoring to use t.Cleanup. This simplifies resource management without changing test behavior.

---

47-50: Simplified cleanup function in test helper

The returned cleanup function no longer calls the database cleanup, as this is now handled internally by testutils.NewDB through t.Cleanup. The function still correctly handles canceling the context and the Anvil cleanup.

pkg/indexer/storer/groupMessage_test.go (2)

23-23: Simplified database initialization

Removed the explicit cleanup function capture from testutils.NewDB, which is now handled internally through t.Cleanup. This change is consistent with the refactoring pattern across the codebase.

---

42-45: Streamlined cleanup function

The cleanup function no longer needs to manage database cleanup as this is handled internally by testutils.NewDB. The function still correctly manages the Anvil cleanup and context cancellation.

pkg/indexer/e2e_test.go (1)

29-29: Database cleanup refactored to use t.Cleanup

This change is part of the broader refactoring to remove explicit cleanup functions from database setup, now relying on Go's t.Cleanup mechanism instead. The database handle is still properly captured while the cleanup function is no longer needed to be explicitly deferred.

pkg/db/congestion_test.go (2)

13-20: Function signature simplified with cleanup handled internally

The setupTest function has been properly refactored to no longer return a cleanup function, aligning with the codebase-wide change to use t.Cleanup for resource management. The function now returns just the context and querier objects.

---

37-37: Test calls updated to match new signature

The test calls have been correctly updated to match the new setupTest signature that no longer returns a cleanup function. This is consistent with the refactoring approach across the codebase.

Also applies to: 57-57, 77-77

pkg/db/subscription_test.go (2)

20-27: Setup function simplified with cleanup handled internally

The setup function signature has been properly updated to no longer return a cleanup function. This is consistent with the project-wide refactoring to use Go's t.Cleanup mechanism in the testutils.NewDB function.

---

122-122: Test calls updated to match new setup signature

All test function calls have been correctly updated to no longer capture or defer the cleanup function, which is now handled automatically via t.Cleanup.

Also applies to: 146-146, 172-172

pkg/db/unsettledUsage_test.go (1)

15-15: Database setup refactored to use t.Cleanup

Direct calls to testutils.NewDB have been updated to no longer capture the cleanup function. This is consistent with the refactoring across the codebase to use Go's t.Cleanup mechanism for resource management instead of explicit cleanup functions.

Also applies to: 57-57

pkg/testutils/store.go (5)

34-42: Improved resource cleanup with t.Cleanup()

The openDB function has been refactored to use t.Cleanup() instead of returning a cleanup function. This is a good practice as it guarantees database connection cleanup even if the test panics, and simplifies the function signature.

---

44-46: Simplified interface by removing cleanup function return

The newCtlDB function now returns only the database and DSN without a cleanup function, making the API more straightforward while still ensuring proper cleanup through the underlying openDB call.

---

48-63: Improved test isolation with guaranteed cleanup

The newInstanceDB function now registers database cleanup with t.Cleanup() instead of returning a cleanup function. This ensures that temporary test databases are always dropped after the test completes, even in case of test failures or panics.

---

65-70: Simplified database setup with implicit cleanup

The NewDB function now has a cleaner signature, returning only the database and DSN. Cleanup is handled implicitly through t.Cleanup() registrations in the called functions, making the API more user-friendly.

---

72-82: Streamlined multiple database creation

The NewDBs function has been simplified to focus on its core responsibility of creating multiple test databases. Cleanup is now handled implicitly by the testing framework, making the code more maintainable.

pkg/registrant/registrant_test.go (9)

39-64: Improved test setup with implicit cleanup

The setup function has been refactored to return only the *deps struct, eliminating the explicit cleanup function. This simplifies the function signature and aligns with Go's testing best practices by relying on t.Cleanup() registered in the underlying NewDB function.

---

66-84: Simplified test helper API

The setupWithRegistrant function now returns only the essential components (*deps and *registrant.Registrant) without the cleanup function. This makes the function easier to use in tests while maintaining proper resource cleanup.

---

87-88: Cleaner test code with implicit cleanup

The test function now calls setup() without needing to capture and defer a cleanup function. This results in more readable and concise test code while still ensuring proper resource cleanup.

---

101-102: Simplified test initialization

Consistent with the cleanup pattern changes, the test now directly uses the returned deps without handling cleanup explicitly.

---

119-121: Clean test initialization without explicit cleanup

The test function initialization is now more straightforward, focusing on the test logic rather than resource management details.

---

167-168: Consistent cleanup pattern application

The test setup is now consistent with the other tests in the file, using the simplified setup() function that implicitly handles cleanup via the testing framework.

---

194-195: Clean test implementation with implicit cleanup

The test function uses the simplified setup() approach consistently with the other tests, making the code more maintainable and readable.

---

221-222: Standardized test initialization

The standardized approach to test setup without explicit cleanup functions maintains consistency throughout the test file.

---

239-240: Simplified setup with registrant

The setupWithRegistrant call now returns only the essential components without an explicit cleanup function, making the test code cleaner and more focused on the actual test logic.

✨ Finishing Touches

• 📝 Generate Docstrings

---

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

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[neekolas]

I've recently discovered t.Context(), which is a context that gets canceled at the end of the test. And it stops you from needing to pass around the context separate from testing.T.

It's a new one in 1.24

[mkysel]

oooh! Love it!