ref(tests): Fix flaky integration test (#314)

Author: iambriccardo

AGENTS.md

@@ -8,6 +8,8 @@
 ## Build and Test
 - Build: `cargo build --workspace --all-targets --all-features`.
 - Lint/format: `cargo fmt`; `cargo clippy --all-targets --all-features -- -D warnings`.
+- Use `ENABLE_TRACING=1` when running integration tests to see the logs.
+- Use `RUST_LOG=[log-level]` if you need to see the logs with a specific log level.
 
 ## Coding Style & Naming
 - Rust 2024 edition; keep formatter clean and warnings denied.

etl/src/test_utils/event.rs

@@ -41,15 +41,32 @@ pub fn group_events_by_type_and_table_id(
 
 pub fn check_events_count(events: &[Event], conditions: Vec<(EventType, u64)>) -> bool {
     let grouped_events = group_events_by_type(events);
-    for (event_type, count) in conditions {
-        let Some(inner_events) = grouped_events.get(&event_type) else {
-            return false;
-        };
 
-        if inner_events.len() != count as usize {
-            return false;
+    conditions.into_iter().all(|(event_type, count)| {
+        grouped_events
+            .get(&event_type)
+            .map(|inner| inner.len() == count as usize)
+            .unwrap_or(false)
+    })
+}
+
+/// Returns a new Vec of events with duplicates removed.
+///
+/// Events that are not tied to a specific row (Begin/Commit/Relation/Truncate/Unsupported)
+/// are not de-duplicated and are preserved in order.
+/// Returns a new Vec of events with duplicates removed based on full equality of events.
+///
+/// Two events are considered the same if all their fields are equal. The first
+/// occurrence is kept and subsequent duplicates are dropped.
+///
+/// The rationale for having this method is that the pipeline doesn't guarantee exactly once delivery
+/// thus in some tests we might have to exclude duplicates while performing assertions.
+pub fn deduplicate_events(events: &[Event]) -> Vec<Event> {
+    let mut result: Vec<Event> = Vec::with_capacity(events.len());
+    for e in events.iter().cloned() {
+        if !result.contains(&e) {
+            result.push(e);
         }
     }
-
-    true
+    result
 }

etl/src/test_utils/test_destination_wrapper.rs

@@ -7,7 +7,7 @@ use tokio::sync::{Notify, RwLock};
 
 use crate::destination::Destination;
 use crate::error::EtlResult;
-use crate::test_utils::event::check_events_count;
+use crate::test_utils::event::{check_events_count, deduplicate_events};
 use crate::types::{Event, EventType, TableRow};
 
 type EventCondition = Box<dyn Fn(&[Event]) -> bool + Send + Sync>;
@@ -100,6 +100,12 @@ impl<D> TestDestinationWrapper<D> {
         self.inner.read().await.events.clone()
     }
 
+    /// Get all events that have been written, de-duplicated by full event equality.
+    pub async fn get_events_deduped(&self) -> Vec<Event> {
+        let events = self.inner.read().await.events.clone();
+        deduplicate_events(&events)
+    }
+
     /// Wait for a specific condition on events
     pub async fn notify_on_events<F>(&self, condition: F) -> Arc<Notify>
     where
@@ -114,12 +120,24 @@ impl<D> TestDestinationWrapper<D> {
         notify
     }
 
-    /// Wait for a specific number of events of given types
+    /// Wait for a specific number of events of given types.
     pub async fn wait_for_events_count(&self, conditions: Vec<(EventType, u64)>) -> Arc<Notify> {
         self.notify_on_events(move |events| check_events_count(events, conditions.clone()))
             .await
     }
 
+    /// Wait for a specific number of events of given types after de-duplicating by full event equality.
+    pub async fn wait_for_events_count_deduped(
+        &self,
+        conditions: Vec<(EventType, u64)>,
+    ) -> Arc<Notify> {
+        self.notify_on_events(move |events| {
+            let deduped = deduplicate_events(events);
+            check_events_count(&deduped, conditions.clone())
+        })
+        .await
+    }
+
     pub async fn clear_events(&self) {
         let mut inner = self.inner.write().await;
         inner.events.clear();

etl/src/types/event.rs

@@ -49,7 +49,7 @@ pub struct CommitEvent {
 pub struct RelationEvent {
     /// LSN position where the event started.
     pub start_lsn: PgLsn,
-    /// LSN position where the event will commit.
+    /// LSN position where the transaction of this event will commit.
     pub commit_lsn: PgLsn,
     /// Complete table schema including columns and types.
     pub table_schema: TableSchema,
@@ -63,7 +63,7 @@ pub struct RelationEvent {
 pub struct InsertEvent {
     /// LSN position where the event started.
     pub start_lsn: PgLsn,
-    /// LSN position where the event will commit.
+    /// LSN position where the transaction of this event will commit.
     pub commit_lsn: PgLsn,
     /// ID of the table where the row was inserted.
     pub table_id: TableId,
@@ -80,7 +80,7 @@ pub struct InsertEvent {
 pub struct UpdateEvent {
     /// LSN position where the event started.
     pub start_lsn: PgLsn,
-    /// LSN position where the event will commit.
+    /// LSN position where the transaction of this event will commit.
     pub commit_lsn: PgLsn,
     /// ID of the table where the row was updated.
     pub table_id: TableId,
@@ -100,11 +100,11 @@ pub struct UpdateEvent {
 /// information about the deleted row for proper cleanup in the destination system.
 #[derive(Debug, Clone, PartialEq)]
 pub struct DeleteEvent {
-    /// LSN position where the event started
+    /// LSN position where the event started.
     pub start_lsn: PgLsn,
-    /// LSN position where the event will commit
+    /// LSN position where the transaction of this event will commit.
     pub commit_lsn: PgLsn,
-    /// ID of the table where the row was deleted
+    /// ID of the table where the row was deleted.
     pub table_id: TableId,
     /// Data from the deleted row.
     ///
@@ -123,7 +123,7 @@ pub struct DeleteEvent {
 pub struct TruncateEvent {
     /// LSN position where the event started.
     pub start_lsn: PgLsn,
-    /// LSN position where the event will commit.
+    /// LSN position where the transaction of this event will commit.
     pub commit_lsn: PgLsn,
     /// Truncate operation options from Postgres.
     pub options: i8,

etl/tests/pipeline.rs

@@ -149,6 +149,15 @@ async fn publication_changes_are_correctly_handled() {
 
     let database = spawn_source_database().await;
 
+    if let Some(server_version) = database.server_version()
+        && server_version.get() <= 150000
+    {
+        println!(
+            "Skipping test for PostgreSQL version <= 15, CREATE PUBLICATION FOR TABLES IN SCHEMA is not supported"
+        );
+        return;
+    }
+
     // Create two tables in the test schema and a publication for that schema.
     let table_1 = test_table_name("table_1");
     let table_1_id = database
@@ -192,15 +201,6 @@ async fn publication_changes_are_correctly_handled() {
     table_1_done.notified().await;
     table_2_done.notified().await;
 
-    if let Some(server_version) = database.server_version()
-        && server_version.get() <= 150000
-    {
-        println!(
-            "Skipping test for PostgreSQL version <= 15, CREATE PUBLICATION FOR TABLES IN SCHEMA is not supported"
-        );
-        return;
-    }
-
     // Insert one row in each table and wait for two insert events.
     let inserts_notify = destination
         .wait_for_events_count(vec![(EventType::Insert, 2)])
@@ -215,10 +215,7 @@ async fn publication_changes_are_correctly_handled() {
         .unwrap();
     inserts_notify.notified().await;
 
-    // Shutdown pipeline before altering publication (by dropping one table)
-    pipeline.shutdown_and_wait().await.unwrap();
-
-    // Drop table_2 so it's no longer part of the publication
+    // Drop table_2 so it's no longer part of the publication.
     database
         .client
         .as_ref()
@@ -230,6 +227,10 @@ async fn publication_changes_are_correctly_handled() {
         .await
         .unwrap();
 
+    // Shutdown pipeline after the table was dropped. We do this to show that the dropping of a table
+    // doesn't cause issues with the pipeline since the change is picked up on pipeline restart.
+    pipeline.shutdown_and_wait().await.unwrap();
+
     // Create table_3 which is going to be added to the publication.
     let table_3 = test_table_name("table_3");
     let table_3_id = database
@@ -258,7 +259,7 @@ async fn publication_changes_are_correctly_handled() {
     // Insert one row in table_1 and wait for it. (We wait for 4 inserts since it keeps the previous
     // ones).
     let inserts_notify = destination
-        .wait_for_events_count(vec![(EventType::Insert, 4)])
+        .wait_for_events_count_deduped(vec![(EventType::Insert, 4)])
         .await;
 
     database
@@ -274,30 +275,32 @@ async fn publication_changes_are_correctly_handled() {
 
     pipeline.shutdown_and_wait().await.unwrap();
 
-    // Assert that table_2 state is gone but destination data remains
+    // Assert that table_2 state is gone but destination data remains.
     let states = store.get_table_replication_states().await;
     assert!(states.contains_key(&table_1_id));
     assert!(!states.contains_key(&table_2_id));
     assert!(states.contains_key(&table_3_id));
 
     // The destination should have the 2 events of the first table, the 1 event of the removed table
     // and the 1 event of the new table.
-    let events = destination.get_events().await;
+    // Use de-duplicated events for assertions to be robust to potential duplicates
+    // on restart where confirmed_flush_lsn may not have been stored.
+    let events = destination.get_events_deduped().await;
     let grouped = group_events_by_type_and_table_id(&events);
     let table_1_inserts = grouped
         .get(&(EventType::Insert, table_1_id))
         .cloned()
-        .unwrap_or_default();
+        .unwrap();
     assert_eq!(table_1_inserts.len(), 2);
     let table_2_inserts = grouped
         .get(&(EventType::Insert, table_2_id))
         .cloned()
-        .unwrap_or_default();
+        .unwrap();
     assert_eq!(table_2_inserts.len(), 1);
     let table_3_inserts = grouped
         .get(&(EventType::Insert, table_3_id))
         .cloned()
-        .unwrap_or_default();
+        .unwrap();
     assert_eq!(table_3_inserts.len(), 1);
 }
 
@@ -307,6 +310,15 @@ async fn publication_for_all_tables_in_schema_ignores_new_tables_until_restart()
 
     let database = spawn_source_database().await;
 
+    if let Some(server_version) = database.server_version()
+        && server_version.get() <= 150000
+    {
+        println!(
+            "Skipping test for PostgreSQL version <= 15, CREATE PUBLICATION FOR TABLES IN SCHEMA is not supported"
+        );
+        return;
+    }
+
     // Create first table.
     let table_1 = test_table_name("table_1");
     let table_1_id = database
@@ -364,15 +376,6 @@ async fn publication_for_all_tables_in_schema_ignores_new_tables_until_restart()
     assert!(table_schemas.contains_key(&table_1_id));
     assert!(!table_schemas.contains_key(&table_2_id));
 
-    if let Some(server_version) = database.server_version()
-        && server_version.get() <= 150000
-    {
-        println!(
-            "Skipping test for PostgreSQL version <= 15, CREATE PUBLICATION FOR TABLES IN SCHEMA is not supported"
-        );
-        return;
-    }
-
     // We restart the pipeline and verify that the new table is now processed.
     let mut pipeline = create_pipeline(
         &database.config,