Kafka test suite

[lquerel]

Pre-filing checklist

• I searched existing issues and didn't find a duplicate

Component(s)

Rust OTAP dataflow (rust/otap-dataflow/)

Objective

NA

Rationale

NA

Scope

NA

Acceptance Criteria

NA

Dependencies or Blockers

No response

Additional Context

No response

[c1ly]

Plan: Kafka Mock Test Suite (common/kafka/testkit)

Goal

Build a reusable, robust test suite that wraps rdkafka's MockCluster so the
Kafka exporter (kafka-exporter branch) and the Kafka receiver
(kafka-receiver branch — already exists) can stand up an in-process mock
broker, connect the component to it, produce/consume messages, inspect broker
state, and drive failure/rebalance scenarios — without Docker. It replaces the
ad-hoc helper functions currently inlined in both kafka_exporter/exporter.rs
and kafka_receiver/receiver.rs and simplifies their integration tests
(conversion need not be 1:1).

Because nearly every integration test also needs a running node, the suite
additionally provides two component wrappers — KafkaExporterHarness and
KafkaReceiverHarness — that spin up the exporter/receiver against the mock
broker and absorb the engine-wiring / spawn / shutdown / ack boilerplate. The
wrappers consume the testkit (they take its KafkaTestCluster); the testkit
core stays broker-only and node-agnostic, so the dependency arrow points one
way: wrapper → testkit + node.

Key facts from research

• rdkafka is pinned to =0.38.0 (workspace Cargo.toml:163), features
  tokio, tracing, plus contrib-nodes adds zstd, ssl (+ platform build
  features). The mocking module is always compiled (no cargo feature needed).
• MockCluster is !Send/!Sync (holds a raw pointer). It must live on
  its creation thread for the broker's whole lifetime; dropping it destroys the
  broker. Current tests therefore run inside a tokio::task::LocalSet.
• Public MockCluster API (actual names in 0.38.0):
  MockCluster::new(broker_count), bootstrap_servers(),
  create_topic(topic, partitions, replication),
  partition_leader(topic, partition, Option<broker_id>),
  partition_follower(...), follower_watermarks(...),
  broker_down(broker_id), broker_up(broker_id) (-1 = all),
  broker_round_trip_time(broker_id, Duration), broker_rack(...),
  coordinator(MockCoordinator::{Group,Transaction}, broker_id),
  apiversion(api_key, min, max), request_errors(api_key, &[RDKafkaRespErr]),
  clear_request_errors(api_key), topic_error(topic, err).
  Note: CreateTopics admin API is unsupported (must use create_topic);
  broker_down/up do not trigger leader election (reassign leaders manually
  via partition_leader); no set_group_final_offset binding in 0.38.0.
• common/kafka/mod.rs is byte-identical on both branches — a safe shared
  home for the testkit.
• rdkafka is an optional normal dependency
  (crates/contrib-nodes/Cargo.toml), pulled in only by the kafka-exporter,
  kafka-receiver, and aws features. It is not a dev-dependency. Both
  branches have identical [dev-dependencies].
• Current exporter test helpers to replace (in
  crates/contrib-nodes/src/exporters/kafka_exporter/exporter.rs):
  start_mock_kafka (1136), create_test_consumer (1150),
  format_header_value (1323), plus the repeated per-test skeleton
  (MockCluster::new(1) → create_topic(_,1,1) → bootstrap_servers() →
  build config → wire_exporter_harness → LocalSet::run_until →
  timeout(30s, consumer.recv()) → assert topic/payload/header/key → shutdown).
• Current receiver test helpers to replace (in
  crates/contrib-nodes/src/receivers/kafka_receiver/receiver.rs, ~4161 lines,
  ~22 #[tokio::test] integration tests incl. 4 multi-consumer rebalance
  tests + a drain test): start_mock_kafka,
  start_mock_kafka_with_partitions(n, topics), create_test_producer (sets
  message.timeout.ms=20000), committed_offset_for (via BaseConsumer +
  TopicPartitionList + committed_offsets → Offset::Offset), plus inline
  "second consumer joins the group, poll until assignment().count() > 0"
  rebalance-trigger loops and pervasive 20×/250ms poll-until-committed loops.
• Payload builders (logs_request_bytes, logs_pdata, traces_pdata,
  metrics_pdata, create_traces_with_spans*, arrow_records_to_bytes) are
  OTLP/OTAP-domain helpers and stay node-side (out of scope).
• Engine wiring stays node-specific — exporter (wire_exporter_harness,
  Harness, spawn_exporter, KeepAlive) and receiver
  (wire_receiver_harness*, setup_receiver_harness*,
  setup_manual_traces_harness*, ack-via-control-channel with
  next_ack/AckMsg) differ. The testkit is Kafka-broker-only and generic
  across both components.
• No state-inspection API on the mock. MockCluster 0.38.0 exposes only
  setters/mutators (create_topic, partition_leader, broker_down/up,
  request_errors, …) plus bootstrap_servers — there is no *_get/query
  binding. Broker/group state is observable only via a normal client probe:
  fetch_metadata (topics/partitions/leaders/replicas/ISR), fetch_watermarks
  (low/high offsets), committed_offsets, and assignment. The testkit
  encapsulates this in a BrokerInspector.
• Component drive surfaces (what the wrappers automate):
  • Exporter (sink): ExporterWrapper::local(KafkaExporter::new(ctx, cfg)?, …)
    → .start(runtime_ctrl_tx, completion_tx, metrics_reporter, Interests::empty()) on the LocalSet; fed pdata via a Sender, controlled
    via control_sender(), ended with NodeControlMsg::Shutdown { deadline }.
    Output lands in the broker (observed via the testkit consumer/inspector).
  • Receiver (source): Box::new(KafkaReceiver::new(ctx, cfg)?) →
    .start(local::ControlChannel, local::EffectHandler); the test reads
    pdata_rx, acks via control_tx + next_ack/AckMsg, optionally
    effect_handler.set_capture_policy(..), ends with Shutdown { deadline }.
    Input comes from the broker (fed via the testkit producer).
  • All wiring types are already available as dev-deps: otap-df-engine
    (test-utils), otap-df-otap (testing), otap-df-pdata (testing),
    otap-df-telemetry. No new dependency is required for the wrappers.

Design decisions (confirmed)

1. Location: new submodule crates/contrib-nodes/src/common/kafka/testkit/
   (in-crate, no new crate). Shared by exporter + receiver via common/kafka.
2. Testkit scope: the testkit core is broker-only and node-agnostic
   (cluster/producer/consumer/message/faults/inspect/wait). It does not
   import node types. Running a component is handled by separate wrappers
   (next decision), not embedded in the testkit.
3. Component wrappers: two role-specific wrappers,
   KafkaExporterHarness (sink) and KafkaReceiverHarness (source), own the
   engine wiring + LocalSet spawn + lifecycle and consume the testkit's
   KafkaTestCluster. They keep channel receiver-ends alive internally (no
   KeepAlive boilerplate in tests) and expose rich handles
   (send / recv / ack / set_capture_policy / shutdown / await).
4. API richness: builders + high-level assertions + broker-state helpers.
5. Threading: LocalSet-based; provide a run/run_on_local_set helper that
   owns the LocalSet and keeps the mock alive for the closure's duration.
   Wrappers spawn the node onto that same LocalSet.
6. Tests: refactor existing exporter AND receiver integration tests to the
   suite + wrappers, AND add capability-demo tests (multi-partition,
   multi-consumer rebalance, committed-offset advancement, broker_down/up
   recovery, rtt, request-error injection, wrapper smoke).
7. Test-only: the testkit and wrappers are a test dependency — never
   compiled into the shipping library and add no new dependency (see next
   section).

Test scoping (no cargo feature, no new dependency)

The testkit is purely test-scoped:

• It lives at common/kafka/testkit/ and is declared #[cfg(test)] in
  common/kafka/mod.rs (#[cfg(test)] mod testkit;, pub(crate) re-exports as
  needed). It therefore compiles only during cargo test, and only when
  a feature that already brings rdkafka is active
  (kafka-exporter / kafka-receiver). It is never part of the shipping lib.
• No new cargo feature (no kafka-test-support) and no new dependency:
  it reuses the rdkafka already present when the node feature is enabled.
  rdkafka stays out of normal (non-test) builds automatically.
• Both the exporter's and the receiver's inline #[cfg(test)] test modules live
  in the same crate (otap-df-contrib-nodes), so they can share this
  crate-internal testkit directly. A #[cfg(test)] module in common/kafka is
  visible to any unit test in the crate, so both nodes' inline tests can
  use crate::common::kafka::testkit::.... (It is intentionally not consumable
  from outside the crate; nothing today needs that.)
• Wrappers live in a #[cfg(test)] module beside the testkit
  (common/kafka/node_harness.rs). Each wrapper is additionally gated by its
  node feature — #[cfg(feature = "kafka-exporter")] for
  KafkaExporterHarness, #[cfg(feature = "kafka-receiver")] for
  KafkaReceiverHarness — so it only compiles when that node (and its
  rdkafka) is present. The testkit core stays node-feature-agnostic; only
  the wrappers reference node types (KafkaExporter/KafkaReceiver, their
  configs) and the engine-wiring dev-deps. No new dependency.
• Branch note: the exporter and receiver are on separate branches. The testkit
  and wrappers land in the shared common/kafka module (identical on both), and
  each branch adopts them in its own test refactor; sequencing/merge is out of
  scope for this plan.

Module layout

common/kafka/
  testkit/          // broker-only, node-agnostic
    mod.rs          // re-exports; KafkaTestCluster; run helpers; prelude
    cluster.rs      // KafkaTestCluster + KafkaTestClusterBuilder (broker/topic setup, state control)
    producer.rs     // TestProducer + builder (multiple producers, keyed/partitioned/header send)
    consumer.rs     // TestConsumer + builder + RebalanceTrigger (groups, offsets, assignment)
    message.rs      // ConsumedMessage view + MessageMatcher + header/key/payload assertions
    faults.rs       // broker-state MUTATORS (down/up, rtt, request_errors, coordinator, leader)
    inspect.rs      // BrokerInspector: read-only topology/watermark/committed-offset queries (probe client)
    wait.rs         // generic poll_until(deadline, interval, predicate) async retry helper
  node_harness.rs   // #[cfg(test)] component WRAPPERS consuming the testkit:
                    //   KafkaExporterHarness  (#[cfg(feature = "kafka-exporter")])
                    //   KafkaReceiverHarness  (#[cfg(feature = "kafka-receiver")])

Public API sketch

KafkaTestCluster (cluster.rs)

Wraps the !Send MockCluster<'static, DefaultProducerContext> and its
bootstrap-servers string.

pub struct KafkaTestCluster { /* mock + brokers */ }

pub struct KafkaTestClusterBuilder {
    broker_count: i32,               // default 1
    topics: Vec<TopicSpec>,          // name, partitions, replication
}

impl KafkaTestClusterBuilder {
    pub fn broker_count(self, n: i32) -> Self;
    pub fn topic(self, name: impl Into<String>) -> Self;              // 1 part / 1 repl
    pub fn topic_with(self, name, partitions: i32, replication: i32) -> Self;
    pub fn build(self) -> KafkaTestCluster;                           // panics w/ ctx on FFI err
}

impl KafkaTestCluster {
    pub fn builder() -> KafkaTestClusterBuilder;
    pub fn new(broker_count: i32) -> Self;                            // convenience
    pub fn bootstrap_servers(&self) -> &str;                          // for component config
    pub fn create_topic(&self, name, partitions: i32, replication: i32);
    pub fn producer(&self) -> TestProducerBuilder;                    // bound to this cluster
    pub fn consumer(&self) -> TestConsumerBuilder;                    // bound to this cluster
    pub fn faults(&self) -> &BrokerFaults;                            // broker-state control
    // raw escape hatch for advanced cases:
    pub fn mock(&self) -> &MockCluster<'static, DefaultProducerContext>;
}

Threading / run helper (mod.rs)

Because MockCluster is !Send, provide a helper that creates a LocalSet,
builds the cluster, and runs the user's async closure to completion with the
mock kept alive:

/// Runs `f` on a current-thread LocalSet with a live cluster.
/// `f` receives the cluster and returns a future.
pub async fn with_cluster<F, Fut, T>(builder: KafkaTestClusterBuilder, f: F) -> T
where
    F: FnOnce(Rc<KafkaTestCluster>) -> Fut,
    Fut: Future<Output = T>;

/// Convenience for the common `#[tokio::test]` shape: sets up LocalSet + cluster.
pub async fn run_on_local_set<F, Fut, T>(f: F) -> T /* builds default 1-broker cluster */;

KafkaTestCluster is exposed as Rc<_> inside the closure so multiple
producers/consumers spawned via spawn_local can share it. The helper wraps the
body in LocalSet::run_until, matching current test structure.

TestProducer (producer.rs) — supports multiple producers

The receiver feeds the broker, so the producer is a co-equal, first-class
utility (not just an exporter afterthought). It must support keys, headers, and
explicit target partitions (the receiver's rebalance tests produce N records to
each specific partition).

pub struct TestProducerBuilder { /* brokers + overrides */ }
impl TestProducerBuilder {
    pub fn client_id(self, id: &str) -> Self;
    pub fn message_timeout(self, d: Duration) -> Self;    // default 20s (matches receiver)
    pub fn set(self, key: &str, val: &str) -> Self;       // raw rdkafka overrides
    pub fn build(self) -> TestProducer;                   // FutureProducer under the hood
}

/// Fluent record spec (payload + optional key/headers/partition).
pub struct SendRecord<'a> { /* topic, payload, key, headers, partition */ }
impl<'a> SendRecord<'a> {
    pub fn new(topic: &'a str, payload: &'a [u8]) -> Self;
    pub fn key(self, key: &'a [u8]) -> Self;
    pub fn partition(self, partition: i32) -> Self;
    pub fn header(self, key: &'a str, value: &'a [u8]) -> Self;   // -> OwnedHeaders
}

impl TestProducer {
    pub async fn send(&self, topic: &str, payload: &[u8]) -> Result<(), TestKitError>;
    pub async fn send_full(&self, rec: SendRecord<'_>) -> Result<(), TestKitError>;
    pub async fn send_to_partition(&self, topic: &str, partition: i32, payload: &[u8])
        -> Result<(), TestKitError>;
    pub async fn send_n(&self, topic: &str, payloads: &[&[u8]]);  // batch
    /// Produce `count` records to each of `partitions`, keyed `k-{p}-{i}`
    /// (the receiver rebalance-test pattern, folded into one call).
    pub async fn produce_per_partition(&self, topic: &str, partitions: i32, count: i32,
        payload: &[u8]);
    pub fn flush(&self, timeout: Duration);
}

TestConsumer (consumer.rs) — multiple consumers/groups, rebalancing, offsets

pub struct TestConsumerBuilder { /* brokers + overrides */ }
impl TestConsumerBuilder {
    pub fn group_id(self, id: &str) -> Self;              // default unique per builder
    pub fn auto_offset_reset(self, v: &str) -> Self;      // default "earliest"
    pub fn enable_auto_commit(self, on: bool) -> Self;    // default false
    pub fn session_timeout(self, d: Duration) -> Self;    // default 6s
    pub fn set(self, key: &str, val: &str) -> Self;
    pub fn subscribe(self, topics: &[&str]) -> TestConsumer;
    pub fn assign_partition(self, topic: &str, partition: i32) -> TestConsumer; // manual assign
}
impl TestConsumer {
    /// Receive one message within `timeout` (default 30s). Double-expect folded in.
    pub async fn recv(&self) -> ConsumedMessage;
    pub async fn try_recv(&self, timeout: Duration) -> Option<ConsumedMessage>;
    /// Consume exactly `n` messages (each bounded by per-msg timeout).
    pub async fn recv_n(&self, n: usize) -> Vec<ConsumedMessage>;
    /// Collect until no message arrives within `idle` (drain semantics).
    pub async fn collect_until_idle(&self, idle: Duration) -> Vec<ConsumedMessage>;
    pub fn assignment(&self) -> Vec<(String, i32)>;       // current partition assignment
    /// Drive recv() until at least `min_partitions` are assigned or timeout.
    pub async fn wait_for_assignment(&self, min_partitions: usize, timeout: Duration) -> bool;

    // ---- consumer-group offset inspection (replaces receiver's committed_offset_for) ----
    /// Committed offset for a group/topic/partition, via an independent probe
    /// client (BaseConsumer + TopicPartitionList + committed_offsets).
    pub fn committed_offset(&self, topic: &str, partition: i32) -> Option<i64>;
    pub fn committed_offsets(&self, tps: &[(&str, i32)]) -> Vec<Option<i64>>;
}

/// A lightweight extra group member used to *force* a rebalance: it joins the
/// same group, is polled until it holds an assignment, and reverts the
/// assignment on drop. Encapsulates the receiver tests' "second consumer joins"
/// pattern (used by 3-4 rebalance tests).
pub struct RebalanceTrigger { /* StreamConsumer */ }
impl RebalanceTrigger {
    pub async fn join(cluster: &KafkaTestCluster, group: &str, topics: &[&str],
        timeout: Duration) -> Self;      // subscribes + polls until assigned
    pub fn assignment(&self) -> Vec<(String, i32)>;
}

A standalone group offset probe is also exposed for callers that don't hold a
TestConsumer:

pub fn committed_offset(cluster: &KafkaTestCluster, group: &str, topic: &str,
    partition: i32) -> Option<i64>;

Multiple consumers sharing a group.id model consumer-group rebalancing;
assignment() / wait_for_assignment() let tests assert partition distribution
after joining, and RebalanceTrigger forces a revoke deterministically.
Multiple producers are just multiple TestProducer instances.

ConsumedMessage + matchers (message.rs) — message checking

Own the borrowed data eagerly (copy out topic/partition/offset/key/payload/
headers) so assertions don't fight rdkafka lifetimes.

pub struct ConsumedMessage {
    pub topic: String, pub partition: i32, pub offset: i64,
    pub key: Option<Vec<u8>>, pub payload: Option<Vec<u8>>,
    pub headers: Vec<(String, Option<Vec<u8>>)>,
}
impl ConsumedMessage {
    pub fn header(&self, key: &str) -> Option<&[u8]>;
    pub fn message_format(&self) -> Option<&[u8]>;        // reads MSG_FORMAT_HEADER
    pub fn assert_topic(&self, t: &str) -> &Self;
    pub fn assert_payload(&self, bytes: &[u8]) -> &Self;
    pub fn assert_key(&self, key: &[u8]) -> &Self;
    pub fn assert_header(&self, key: &str, val: &[u8]) -> &Self;
    pub fn assert_format_otlp(&self) -> &Self;            // == MSG_FORMAT_OTLP
    pub fn assert_format_otap(&self) -> &Self;            // == MSG_FORMAT_OTAP
}

// Builder-style matcher for collections:
pub struct MessageMatcher { /* topic/key/header/payload predicates */ }
pub fn count_by_topic(msgs: &[ConsumedMessage]) -> HashMap<String, usize>;
pub fn count_by_partition(msgs: &[ConsumedMessage]) -> HashMap<(String,i32), usize>;

This replaces format_header_value and the repeated topic/payload/header
assertion trio.

BrokerFaults (faults.rs) — broker state + fault injection

Thin, well-documented wrappers over the raw MockCluster state methods, with
doc notes about the 0.38.0 caveats (no auto leader election, etc.).

pub struct BrokerFaults<'a>(&'a MockCluster<'static, DefaultProducerContext>);
impl BrokerFaults<'_> {
    pub fn broker_down(&self, broker_id: i32);            // -1 = all
    pub fn broker_up(&self, broker_id: i32);
    pub fn all_brokers_down(&self);                       // broker_down(-1)
    pub fn all_brokers_up(&self);
    pub fn round_trip_time(&self, broker_id: i32, d: Duration);
    pub fn set_partition_leader(&self, topic: &str, partition: i32, broker: Option<i32>);
    pub fn set_group_coordinator(&self, group: &str, broker: i32);
    pub fn inject_request_errors(&self, api: RDKafkaApiKey, errors: &[RDKafkaRespErr]);
    pub fn clear_request_errors(&self, api: RDKafkaApiKey);
    pub fn set_topic_error(&self, topic: &str, err: RDKafkaRespErr);
    pub fn set_api_version(&self, api: RDKafkaApiKey, min: Option<i16>, max: Option<i16>);
}

The faults helpers are all mutators. Read-only broker/group-state
inspection is a first-class, separate API (inspect.rs) — kept separate
because MockCluster exposes no getters, so inspection must go through a normal
client probe.

BrokerInspector (inspect.rs) — read-only broker/group state

Uses a short-lived BaseConsumer probe bound to the cluster. Topology comes
from fetch_metadata, offsets from fetch_watermarks, and committed group
offsets from committed_offsets (group-scoped probe).

impl KafkaTestCluster {
    pub fn inspect(&self) -> BrokerInspector;                    // topology + watermarks
    pub fn inspect_group(&self, group: &str) -> BrokerInspector; // + committed offsets
}

pub struct BrokerInspector { /* probe BaseConsumer + optional group + probe timeout */ }
impl BrokerInspector {
    pub fn probe_timeout(self, d: Duration) -> Self;             // default 10s, configurable

    // Topology (fetch_metadata)
    pub fn topics(&self) -> Vec<TopicInfo>;                      // name, partition count
    pub fn partitions(&self, topic: &str) -> Vec<PartitionInfo>; // id, leader, replicas, isr
    pub fn topic_exists(&self, topic: &str) -> bool;

    // Offsets / watermarks (fetch_watermarks)
    pub fn watermarks(&self, topic: &str, partition: i32) -> (i64, i64); // (low, high)
    pub fn message_count(&self, topic: &str, partition: i32) -> i64;     // high - low

    // Group state (committed_offsets; requires inspect_group)
    pub fn committed_offset(&self, topic: &str, partition: i32) -> Option<i64>;
    pub fn committed_offsets(&self, tps: &[(&str, i32)]) -> Vec<Option<i64>>;

    pub fn snapshot(&self) -> BrokerStateSnapshot;               // topics+partitions+leaders+watermarks
}

pub struct BrokerStateSnapshot { /* topics, partitions, leaders, replicas, watermarks */ }

TestConsumer::committed_offset (above) remains a convenience for an active
consumer; BrokerInspector is the standalone / group-scoped inspection path and
the home for topology + watermark queries.

Component wrappers (node_harness.rs) — spin up the exporter/receiver

Two role-specific wrappers own all engine wiring, spawn the node onto the
LocalSet that with_cluster already owns, keep channel receiver-ends alive
internally, and expose intention-revealing handles. They consume the
testkit: start takes &KafkaTestCluster and auto-sets bootstrap.servers
from cluster.bootstrap_servers(). Config builders remain fully overridable for
signal/format/commit/rebalance options; start_for provides a sensible default
config for the given topics.

// #[cfg(feature = "kafka-exporter")] — sink: output lands in the broker
pub struct KafkaExporterHarness { /* ExporterWrapper + pdata/control/rt/completion txs + join handle */ }
impl KafkaExporterHarness {
    pub fn start(cluster: &KafkaTestCluster, cfg: KafkaExporterConfig) -> Self;
    pub fn start_for(cluster: &KafkaTestCluster, topics: KafkaTopics) -> Self; // default cfg
    pub async fn send_pdata(&self, pdata: OtapPdata) -> Result<(), TestKitError>;
    pub async fn shutdown(&self, deadline: Duration);   // NodeControlMsg::Shutdown
    pub async fn await_stopped(self);                   // join the spawned task
}

// #[cfg(feature = "kafka-receiver")] — source: input comes from the broker
pub struct KafkaReceiverHarness { /* Box<KafkaReceiver> + ctrl/pdata channels + EffectHandler + join handle */ }
impl KafkaReceiverHarness {
    pub fn start(cluster: &KafkaTestCluster, cfg: KafkaReceiverConfig) -> Self;
    pub fn start_for(cluster: &KafkaTestCluster, topics: KafkaTopics) -> Self;
    pub fn set_capture_policy(&mut self, policy: Option<HeaderCapturePolicy>);
    pub async fn recv_pdata(&mut self) -> OtapPdata;                    // default 30s timeout
    pub async fn try_recv_pdata(&mut self, timeout: Duration) -> Option<OtapPdata>;
    pub fn ack(&self, pdata: OtapPdata);                               // folds next_ack / AckMsg
    pub async fn shutdown(&self, deadline: Duration);
    pub async fn await_stopped(self);
}

KafkaTopics is a small helper describing the per-signal topic/format layout
(traces/metrics/logs) used by both start_for variants. The wrappers do not
build OTLP/OTAP payloads — tests pass OtapPdata in (exporter) or read it out
(receiver), keeping payload-domain concerns in the node test modules.

poll_until (wait.rs) — async retry helper

Both the receiver rebalance tests and any broker-state test repeatedly hand-roll
"loop up to N times, sleep M ms, check condition" (e.g. poll 20× / 250ms until
a committed offset advances, or until a consumer gets an assignment). Provide
one helper:

/// Polls `predicate` until it returns true or `timeout` elapses, sleeping
/// `interval` between attempts. Returns whether the predicate succeeded.
pub async fn poll_until<F>(timeout: Duration, interval: Duration, predicate: F) -> bool
where F: FnMut() -> bool;

/// Async-predicate variant (predicate returns a future), for cases that must
/// drive `recv()` or query the broker between attempts.
pub async fn poll_until_async<F, Fut>(timeout: Duration, interval: Duration, predicate: F) -> bool
where F: FnMut() -> Fut, Fut: Future<Output = bool>;

Error type

TestKitError (small enum: ClusterSetup, Produce, Consume, Timeout)
with Display. Helpers that are clearly test-only may expect(..) with rich
context instead of returning errors, matching current style; send/consume return
Result so negative tests can assert failures.

Configurable timeouts

All timeouts have defaults (matching the current tests) and are overridable:

Knob	Default	Set via
producer message.timeout.ms	20s	TestProducerBuilder::message_timeout(d)
consumer recv	30s	per-call arg / try_recv(timeout)
consumer session.timeout.ms	6s	TestConsumerBuilder::session_timeout(d)
wait_for_assignment / poll_until	—	explicit timeout + interval args
BrokerInspector probe	10s	BrokerInspector::probe_timeout(d)

Refactor of existing exporter tests

Convert the MockCluster-backed tests in exporter.rs (T1–T7:
exports_logs/traces/metrics_otlp_to_mock_broker,
exports_logs_otap_sets_otap_format_header,
routes_to_topic_from_transport_header,
sets_partition_key_from_transport_headers,
shutdown_flushes_buffered_records) to use the testkit:

• Replace start_mock_kafka + create_test_consumer + manual LocalSet +
  timeout(30s, recv()) + format_header_value with
  run_on_local_set / with_cluster, cluster.consumer()...,
  consumer.recv().await, and ConsumedMessage::assert_*.
• Replace wire_exporter_harness / Harness / spawn_exporter /
  send_shutdown with KafkaExporterHarness::start(&cluster, cfg) +
  send_pdata(..) / shutdown(..). The wrapper owns all engine wiring.
• Keep OTLP payload builders (logs_pdata, traces_pdata, metrics_pdata) in
  exporter.rs; they build the OtapPdata handed to send_pdata.
• Delete the now-redundant start_mock_kafka, create_test_consumer,
  format_header_value, and the local harness/spawn/shutdown helpers.
• Conversion is not 1:1: e.g. shutdown_flushes_buffered_records uses
  consumer.recv_n(3); multi-topic routing uses two cluster.topic(...) +
  a single consumer subscribed to the dynamic topic.

Refactor of existing receiver tests (kafka-receiver branch)

The receiver's inline test module already contains the seed helpers this testkit
generalizes. Convert its ~22 integration tests to the testkit:

• Replace start_mock_kafka / start_mock_kafka_with_partitions(n, topics) with
  KafkaTestCluster::builder().topic_with(name, n, 1)...build() (via the run
  helper).
• Replace create_test_producer + FutureRecord::to().payload().key() .partition().headers() with TestProducer / SendRecord /
  produce_per_partition.
• Replace committed_offset_for with committed_offset(...) /
  inspect_group(g).committed_offset(...).
• Replace the inline "second consumer joins the group, poll until
  assignment().count() > 0" blocks with RebalanceTrigger::join(...).
• Replace the pervasive 20×/250ms poll-until-committed loops with poll_until.
• Replace wire_receiver_harness* / setup_receiver_harness* /
  setup_manual_traces_harness* + inline spawn/ack/shutdown with
  KafkaReceiverHarness::start(&cluster, cfg) + recv_pdata() / ack(..) /
  set_capture_policy(..) / shutdown(..). The wrapper owns all engine wiring
  and folds next_ack/AckMsg.
• Keep OTLP/OTAP payload builders and payload→OTLP decode helpers in
  receiver.rs; the test asserts on the OtapPdata returned by recv_pdata.

Tests to convert (grouped): 8 signal/format round-trips (..._traces/logs/ metrics + _otap), 4 header-capture/extraction tests, 4 rebalance tests
(rebalance_single_consumer_assigns_and_commits,
rebalance_revoke_commits_before_reassign,
rebalance_cooperative_sticky_retains_owned_partitions,
rebalance_revoke_then_reassign_preserves_new_records), and
drain_ingress_stops_polling_and_notifies_drained.

New capability-demo tests

Add a testkit-level test module (in testkit/mod.rs under #[cfg(test)]) that
exercises the suite directly with a plain producer/consumer (no exporter), to
prove robustness and document usage:

1. multi-partition produce/consume: topic with 3 partitions; produce_per_ partition; assert count_by_partition distribution and full round-trip.
2. multiple consumers / rebalance: two consumers in one group on a
   3-partition topic (or a RebalanceTrigger); use wait_for_assignment and
   assert assignment() splits partitions across them.
3. committed-offset advancement: consume + commit; poll_until an
   inspect_group(g).committed_offset(...) reaches the expected watermark.
4. topology / watermark inspection: assert inspect().partitions(t).len()
   and inspect().message_count(t, p) after producing.
5. produce-to-specific-partition: send_to_partition; assert the consumed
   message lands on that partition.
6. multiple producers: two producers to the same topic; assert all messages
   consumed.
7. broker_down/up recovery: produce, all_brokers_down(), attempt produce
   (expect failure/retry), all_brokers_up(), assert delivery resumes.
8. rtt latency: set round_trip_time and assert a produce still succeeds
   within an extended timeout (sanity of the knob).
9. request-error injection: inject_request_errors(Metadata, [..]) then
   clear_request_errors; assert the client recovers.
10. wrapper smoke (feature-gated): KafkaExporterHarness::start_for →
    send_pdata → consume via the testkit consumer; and
    KafkaReceiverHarness::start_for ← feed via the testkit producer →
    recv_pdata + ack. Validates the wrappers end-to-end.

Demos 1–9 live with the testkit (broker-only, no node feature required). Demo 10
lives with the wrappers and is gated by kafka-exporter / kafka-receiver.
These validate the suite itself.

Expected usage examples

Illustrations of what a test is expected to look like once the testkit lands.
Snippets are indicative (names may shift in implementation); the node-specific
engine wiring is elided as /* ... */.

Exporter round-trip. Spin up the exporter via its wrapper, send pdata,
assert the produced record.

with_cluster(KafkaTestCluster::builder().topic("it-logs"), |cluster| async move {
    let consumer = cluster.consumer().subscribe(&["it-logs"]);
    let exporter = KafkaExporterHarness::start_for(&cluster, topics);
    exporter.send_pdata(logs_pdata(bytes, None)).await?;
    consumer.recv().await
        .assert_topic("it-logs")
        .assert_payload(&expected_bytes)
        .assert_format_otlp();
    exporter.shutdown(secs(1)).await;
}).await;

Receiver feed + consume. Produce into a multi-partition topic via the
producer, spin up the receiver via its wrapper, read + ack decoded pdata.

with_cluster(KafkaTestCluster::builder().topic_with("t", 2, 1), |cluster| async move {
    let producer = cluster.producer().build();
    producer.produce_per_partition("t", 2, 5, &otlp_bytes).await;
    let mut receiver = KafkaReceiverHarness::start_for(&cluster, topics);
    let pdata = receiver.recv_pdata().await;   // assert on pdata
    receiver.ack(pdata);
    receiver.shutdown(secs(5)).await;
}).await;

Broker-state inspection. Read topology, watermarks, and committed offsets.

let insp = cluster.inspect();
assert_eq!(insp.partitions("t").len(), 3);
assert_eq!(insp.message_count("t", 0), 5);
assert!(cluster.inspect_group("g").committed_offset("t", 0).is_some());

Fault injection + recovery.

cluster.faults().all_brokers_down();
/* produce attempts fail / retry */
cluster.faults().all_brokers_up();
assert!(poll_until_async(secs(10), ms(250),
    || async { producer.send("t", &b).await.is_ok() }).await);

Rebalance (before → after). This is where the boilerplate reduction is
largest, so it is shown as a contrast; all other examples are after-only.

Before (receiver branch today, ~25 lines):

// second consumer joins the SAME group to force a revoke
let consumer_b: StreamConsumer = ClientConfig::new()
    .set("bootstrap.servers", &brokers)
    .set("group.id", group)
    .set("enable.auto.commit", "false")
    .set("auto.offset.reset", "earliest")
    .create().expect("create consumer B");
consumer_b.subscribe(&[topic]).expect("subscribe");
let mut assigned_b = false;
for _ in 0..40 {
    if let Ok(a) = consumer_b.assignment() { if a.count() > 0 { assigned_b = true; break; } }
    let _ = timeout(Duration::from_millis(500), consumer_b.recv()).await;
}
assert!(assigned_b, "rebalance did not occur");
let mut all_committed = false;
for _ in 0..20 {
    let c0 = committed_offset_for(&brokers, group, topic, 0);
    let c1 = committed_offset_for(&brokers, group, topic, 1);
    if c0.is_some_and(|o| o >= 5) && c1.is_some_and(|o| o >= 5) { all_committed = true; break; }
    tokio::time::sleep(Duration::from_millis(250)).await;
}
assert!(all_committed, "commit-before-revoke violated");

After (testkit + wrapper):

let mut receiver = KafkaReceiverHarness::start(&cluster, manual_commit_cfg);
// consume + ack A's records so its offsets become committable ...
let _trigger = RebalanceTrigger::join(&cluster, group, &["t"], secs(10)).await; // forces revoke
let insp = cluster.inspect_group(group);
assert!(poll_until(secs(5), ms(250), || {
    insp.committed_offset("t", 0).is_some_and(|o| o >= 5)
        && insp.committed_offset("t", 1).is_some_and(|o| o >= 5)
}).await, "commit-before-revoke violated");
receiver.shutdown(secs(5)).await;

Implementation steps

1. Add #[cfg(test)] mod testkit; and #[cfg(test)] mod node_harness; in
   common/kafka/mod.rs (no feature, no dependency change).
2. Implement cluster.rs (KafkaTestCluster + builder + create_topic +
   bootstrap_servers + mock() escape hatch).
3. Implement mod.rs run helpers (with_cluster, run_on_local_set) with the
   Rc<KafkaTestCluster> + LocalSet pattern, and wait.rs (poll_until /
   poll_until_async).
4. Implement producer.rs (keyed/partitioned/header send + produce_per_ partition), consumer.rs (offset probes, wait_for_assignment,
   RebalanceTrigger), message.rs, faults.rs (mutators), and inspect.rs
   (BrokerInspector: topology, watermarks, committed offsets via probe client).
5. Implement node_harness.rs wrappers consuming the testkit + node types:
   KafkaExporterHarness (#[cfg(feature = "kafka-exporter")]) and
   KafkaReceiverHarness (#[cfg(feature = "kafka-receiver")]) — build config
   from cluster.bootstrap_servers(), own the engine wiring + LocalSet spawn +
   lifecycle, expose the rich handles.
6. Add the testkit-level capability-demo tests + the wrapper smoke test.
7. Refactor exporter integration tests (T1–T7) onto the testkit + wrapper;
   remove dead exporter helpers (start_mock_kafka, create_test_consumer,
   format_header_value, wire_exporter_harness, Harness, spawn_exporter,
   send_shutdown).
8. Refactor receiver integration tests onto the testkit + wrapper; remove dead
   receiver helpers (start_mock_kafka, start_mock_kafka_with_partitions,
   create_test_producer, committed_offset_for, wire_receiver_harness*,
   setup_receiver_harness*, setup_manual_traces_harness*, inline
   rebalance/poll loops).
9. Run, per branch (each node feature brings rdkafka and enables its wrapper):
   • exporter branch: cd rust/otap-dataflow && cargo check -p otap-df-contrib-nodes --features kafka-exporter, then cargo test -p otap-df-contrib-nodes --features kafka-exporter.
   • receiver branch: same with --features kafka-receiver.
     Keep all MockCluster + node usage inside a LocalSet (!Send). Finally
     cargo xtask check.
10. Changelog: internal test infrastructure is not user-facing → mark PR title
    chore and skip the changelog per AGENTS.md.

Resolved decisions

• Broker-state inspection: kept and promoted to a first-class
  BrokerInspector (topology, watermarks, committed offsets). Since
  MockCluster exposes no getters, inspection goes through a normal client
  probe (fetch_metadata / fetch_watermarks / committed_offsets).
• Timeouts: configurable, with defaults matching the current tests (producer
  20s, consumer recv 30s, session 6s, inspector probe 10s; poll_until /
  wait_for_assignment take explicit timeout + interval).
• Component spin-up: handled by two role-specific wrappers
  (KafkaExporterHarness, KafkaReceiverHarness) in a #[cfg(test)]
node_harness.rs beside the testkit, each gated by its node feature. They
  consume KafkaTestCluster (testkit stays broker-only / node-agnostic),
  own the engine wiring + spawn + lifecycle, and expose rich handles
  (send / recv / ack / set_capture_policy / shutdown / await).