test(connectors): add restart integration test for sink connector

Signed-off-by: shin <sars21@hanmail.net>

Author: seokjin0414

core/integration/tests/connectors/postgres/mod.rs

@@ -19,6 +19,7 @@
 
 mod postgres_sink;
 mod postgres_source;
+mod restart;
 
 use crate::connectors::TestMessage;
 use serde::Deserialize;

core/integration/tests/connectors/postgres/restart.rs

@@ -0,0 +1,163 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+use super::{POLL_ATTEMPTS, POLL_INTERVAL_MS, TEST_MESSAGE_COUNT};
+use crate::connectors::fixtures::{PostgresOps, PostgresSinkFixture};
+use crate::connectors::{TestMessage, create_test_messages};
+use bytes::Bytes;
+use iggy::prelude::{IggyMessage, Partitioning};
+use iggy_binary_protocol::MessageClient;
+use iggy_common::Identifier;
+use iggy_connector_sdk::api::{ConnectorStatus, SinkInfoResponse};
+use integration::harness::seeds;
+use integration::iggy_harness;
+use reqwest::Client;
+use std::time::Duration;
+use tokio::time::sleep;
+
+const API_KEY: &str = "test-api-key";
+const SINK_TABLE: &str = "iggy_messages";
+const SINK_KEY: &str = "postgres";
+
+type SinkRow = (i64, String, String, Vec<u8>);
+
+async fn wait_for_sink_status(
+    http: &Client,
+    api_url: &str,
+    expected: ConnectorStatus,
+) -> SinkInfoResponse {
+    for _ in 0..POLL_ATTEMPTS {
+        if let Ok(resp) = http
+            .get(format!("{api_url}/sinks/{SINK_KEY}"))
+            .header("api-key", API_KEY)
+            .send()
+            .await
+            && let Ok(info) = resp.json::<SinkInfoResponse>().await
+            && info.status == expected
+        {
+            return info;
+        }
+        sleep(Duration::from_millis(POLL_INTERVAL_MS)).await;
+    }
+    panic!("Sink connector did not reach {expected:?} status in time");
+}
+
+fn build_messages(messages_data: &[TestMessage], id_offset: usize) -> Vec<IggyMessage> {
+    messages_data
+        .iter()
+        .enumerate()
+        .map(|(i, msg)| {
+            let payload = serde_json::to_vec(msg).expect("Failed to serialize message");
+            IggyMessage::builder()
+                .id((id_offset + i + 1) as u128)
+                .payload(Bytes::from(payload))
+                .build()
+                .expect("Failed to build message")
+        })
+        .collect()
+}
+
+#[iggy_harness(
+    server(connectors_runtime(config_path = "tests/connectors/postgres/sink.toml")),
+    seed = seeds::connector_stream
+)]
+async fn restart_sink_connector_continues_processing(
+    harness: &TestHarness,
+    fixture: PostgresSinkFixture,
+) {
+    let client = harness.root_client().await.unwrap();
+    let api_url = harness
+        .connectors_runtime()
+        .expect("connector runtime should be available")
+        .http_url();
+    let http = Client::new();
+    let pool = fixture.create_pool().await.expect("Failed to create pool");
+
+    fixture.wait_for_table(&pool, SINK_TABLE).await;
+
+    let stream_id: Identifier = seeds::names::STREAM.try_into().unwrap();
+    let topic_id: Identifier = seeds::names::TOPIC.try_into().unwrap();
+
+    wait_for_sink_status(&http, &api_url, ConnectorStatus::Running).await;
+
+    let first_batch = create_test_messages(TEST_MESSAGE_COUNT);
+    let mut messages = build_messages(&first_batch, 0);
+    client
+        .send_messages(
+            &stream_id,
+            &topic_id,
+            &Partitioning::partition_id(0),
+            &mut messages,
+        )
+        .await
+        .expect("Failed to send first batch");
+
+    let query = format!(
+        "SELECT iggy_offset, iggy_stream, iggy_topic, payload FROM {SINK_TABLE} ORDER BY iggy_offset"
+    );
+    let rows: Vec<SinkRow> = fixture
+        .fetch_rows_as(&pool, &query, TEST_MESSAGE_COUNT)
+        .await
+        .expect("Failed to fetch first batch rows");
+
+    assert_eq!(
+        rows.len(),
+        TEST_MESSAGE_COUNT,
+        "Expected {TEST_MESSAGE_COUNT} rows before restart"
+    );
+
+    let resp = http
+        .post(format!("{api_url}/sinks/{SINK_KEY}/restart"))
+        .header("api-key", API_KEY)
+        .send()
+        .await
+        .expect("Failed to call restart endpoint");
+
+    assert_eq!(
+        resp.status().as_u16(),
+        204,
+        "Restart endpoint should return 204 No Content"
+    );
+
+    wait_for_sink_status(&http, &api_url, ConnectorStatus::Running).await;
+
+    let second_batch = create_test_messages(TEST_MESSAGE_COUNT);
+    let mut messages = build_messages(&second_batch, TEST_MESSAGE_COUNT);
+    client
+        .send_messages(
+            &stream_id,
+            &topic_id,
+            &Partitioning::partition_id(0),
+            &mut messages,
+        )
+        .await
+        .expect("Failed to send second batch");
+
+    let total_expected = TEST_MESSAGE_COUNT * 2;
+    let rows: Vec<SinkRow> = fixture
+        .fetch_rows_as(&pool, &query, total_expected)
+        .await
+        .expect("Failed to fetch rows after restart");
+
+    assert!(
+        rows.len() >= total_expected,
+        "Expected at least {total_expected} rows after restart, got {}",
+        rows.len()
+    );
+}