diff --git a/sqlflow/sources/base.py b/sqlflow/sources/base.py
index b6eee92..656b20c 100644
--- a/sqlflow/sources/base.py
+++ b/sqlflow/sources/base.py
@@ -6,7 +6,7 @@
 logger = logging.getLogger(__name__)
 
 class Message:
-    def __init__(self, value: bytes, topic: str | None, partition: int | None, offset: int | None):
+    def __init__(self, value: bytes, topic: str | None = None, partition: int | None = None, offset: int | None = None):
         self._value = value
         self._topic = topic
         self._partition = partition
diff --git a/tests/integration/test_integration.py b/tests/integration/test_integration.py
index 470d33c..4293da6 100644
--- a/tests/integration/test_integration.py
+++ b/tests/integration/test_integration.py
@@ -384,6 +384,6 @@ def test_dlq_functionality_handler_invoke(bootstrap_server):
     dlq_messages = read_all_kafka_messages(bootstrap_server, dlq_topic)
     assert len(dlq_messages) == 1, f"Expected 1 DLQ message, but got {len(dlq_messages)}"
     m = dlq_messages[0]
-    assert m['error'] == 'Binder Error: Referenced column "broken" not found in FROM clause!\nCandidate bindings: "valid"\n\nLINE 2:   broken\n          ^'
+    assert 'Binder Error: Referenced column "broken" not found in FROM clause!' in m['error']
     assert m['message'] == 'Handler invocation failed'
     assert m['phase'] == 'handler.invoke'
diff --git a/tests/release/test_image.py b/tests/release/test_image.py
index 2c34dd4..5f4b3e3 100644
--- a/tests/release/test_image.py
+++ b/tests/release/test_image.py
@@ -92,5 +92,5 @@ def test_basic_agg_mem_readme_example(git_sha):
     )
 
     messages = read_all_kafka_messages(bootstrap_server, out_topic)
-    assert 1000 == len(messages)
+    assert 5 == len(messages)