iAs/ias-common/tests/jetstream_integration.rs

use ias_common::mq::{
    ConsumerSpec, DeadLetterPayload, InboundPayload, JetStreamBus, JetStreamConfig,
    MessageEnvelope, StreamKind, StreamNames, inbound_subject,
};

#[tokio::test]
#[ignore = "requires a running NATS JetStream server from docker-compose.nats.yml"]
async fn jetstream_publish_and_ack_roundtrip() -> anyhow::Result<()> {
    let config = JetStreamConfig::from_env();
    let bus = JetStreamBus::connect(config.clone()).await?;
    bus.ensure_streams().await?;
    let subject = inbound_subject("wechat", "integration");
    let payload = InboundPayload {
        channel: "wechat".to_string(),
        account_id: "integration".to_string(),
        from_user_id: "tester".to_string(),
        content: "ping".to_string(),
        context: None,
        external_message_id: Some("integration-msg-1".to_string()),
        session_id: None,
        group_id: None,
        create_time_ms: None,
    };
    let envelope = MessageEnvelope::new(
        subject.as_str(),
        payload.idempotency_key(),
        "trace-integration",
        payload,
    );
    bus.publish_envelope(subject.as_str(), &envelope).await?;

    let names = StreamNames::new(config.stream_prefix.clone());
    let spec = ConsumerSpec {
        stream: names.name(StreamKind::Inbound),
        durable_name: "integration-ack-test".to_string(),
        filter_subject: subject.to_string(),
        ack_wait: config.ack_wait(),
        max_deliver: config.max_deliver,
    };
    let (tx, mut rx) = tokio::sync::mpsc::channel(1);
    let consumer_bus = bus.clone();
    let handle = tokio::spawn(async move {
        consumer_bus
            .consume_envelopes::<InboundPayload, _, _>(spec, move |message| {
                let tx = tx.clone();
                async move {
                    tx.send(message.message_id).await?;
                    Ok(())
                }
            })
            .await
    });

    let message_id = tokio::time::timeout(std::time::Duration::from_secs(10), rx.recv())
        .await?
        .ok_or_else(|| anyhow::anyhow!("consumer closed before receiving test message"))?;
    handle.abort();

    assert_eq!(message_id, "wechat/integration/integration-msg-1");
    Ok(())
}

#[tokio::test]
#[ignore = "requires a running NATS JetStream server from docker-compose.nats.yml"]
async fn jetstream_redelivery_reaches_dlq_after_max_deliver() -> anyhow::Result<()> {
    let mut config = JetStreamConfig::from_env();
    config.ack_wait_secs = 1;
    config.max_deliver = 2;
    let bus = JetStreamBus::connect(config.clone()).await?;
    bus.ensure_streams().await?;
    let names = StreamNames::new(config.stream_prefix.clone());
    let unique = uuid::Uuid::new_v4().to_string();
    let account_id = format!("dlq-{unique}");
    let subject = inbound_subject("wechat", &account_id);
    let payload = InboundPayload {
        channel: "wechat".to_string(),
        account_id: account_id.clone(),
        from_user_id: "tester".to_string(),
        content: "force retry".to_string(),
        context: None,
        external_message_id: Some(unique.clone()),
        session_id: None,
        group_id: None,
        create_time_ms: None,
    };
    let message_id = format!("wechat/{account_id}/{unique}");
    let envelope = MessageEnvelope::new(
        subject.as_str(),
        message_id.clone(),
        format!("trace-{unique}"),
        payload,
    );

    let (dlq_tx, mut dlq_rx) = tokio::sync::mpsc::channel(1);
    let dlq_bus = bus.clone();
    let dlq_stream = names.name(StreamKind::Dlq);
    let dlq_ack_wait = config.ack_wait();
    let dlq_max_deliver = config.max_deliver;
    let dlq_handle = tokio::spawn(async move {
        dlq_bus
            .consume_envelopes::<DeadLetterPayload, _, _>(
                ConsumerSpec {
                    stream: dlq_stream,
                    durable_name: format!("integration-dlq-{unique}"),
                    filter_subject: "dlq.>".to_string(),
                    ack_wait: dlq_ack_wait,
                    max_deliver: dlq_max_deliver,
                },
                move |message| {
                    let dlq_tx = dlq_tx.clone();
                    let expected = message_id.clone();
                    async move {
                        if message.payload.original_message_id.as_deref() == Some(&expected) {
                            dlq_tx.send(message.payload.original_message_id).await?;
                        }
                        Ok(())
                    }
                },
            )
            .await
    });

    let retry_bus = bus.clone();
    let retry_spec = ConsumerSpec {
        stream: StreamNames::new(config.stream_prefix.clone()).name(StreamKind::Inbound),
        durable_name: format!("integration-redelivery-{account_id}"),
        filter_subject: subject.to_string(),
        ack_wait: config.ack_wait(),
        max_deliver: config.max_deliver,
    };
    let retry_handle = tokio::spawn(async move {
        retry_bus
            .consume_envelopes::<InboundPayload, _, _>(retry_spec, |_message| async move {
                anyhow::bail!("force redelivery")
            })
            .await
    });

    bus.publish_envelope(subject.as_str(), &envelope).await?;
    let observed = tokio::time::timeout(std::time::Duration::from_secs(20), dlq_rx.recv())
        .await?
        .ok_or_else(|| anyhow::anyhow!("dlq consumer closed before receiving test message"))?;
    retry_handle.abort();
    dlq_handle.abort();

    assert_eq!(observed.as_deref(), Some(envelope.message_id.as_str()));
    Ok(())
}

#[tokio::test]
#[ignore = "requires a running NATS JetStream server from docker-compose.nats.yml"]
async fn jetstream_concurrent_dlq_does_not_lose_messages() -> anyhow::Result<()> {
    let mut config = JetStreamConfig::from_env();
    config.ack_wait_secs = 1;
    config.max_deliver = 2;
    config.consumer_concurrency = 4;
    let bus = JetStreamBus::connect(config.clone()).await?;
    bus.ensure_streams().await?;
    let names = StreamNames::new(config.stream_prefix.clone());
    let batch = uuid::Uuid::new_v4().to_string();
    let count = 8usize;
    let account_id = format!("cdlq-{batch}");

    // 发布 count 条不同 message_id 的消息
    let mut message_ids = Vec::new();
    for index in 0..count {
        let external = format!("{batch}-{index}");
        let subject = inbound_subject("wechat", &account_id);
        let payload = InboundPayload {
            channel: "wechat".to_string(),
            account_id: account_id.clone(),
            from_user_id: "tester".to_string(),
            content: "force retry".to_string(),
            context: None,
            external_message_id: Some(external.clone()),
            session_id: None,
            group_id: None,
            create_time_ms: None,
        };
        let message_id = format!("wechat/{account_id}/{external}");
        message_ids.push(message_id.clone());
        let envelope = MessageEnvelope::new(
            subject.as_str(),
            message_id,
            format!("trace-{batch}-{index}"),
            payload,
        );
        bus.publish_envelope(subject.as_str(), &envelope).await?;
    }

    // DLQ consumer 收集 original_message_id
    let (dlq_tx, mut dlq_rx) = tokio::sync::mpsc::channel(count);
    let dlq_bus = bus.clone();
    let dlq_stream = names.name(StreamKind::Dlq);
    let dlq_ack_wait = config.ack_wait();
    let dlq_max_deliver = config.max_deliver;
    let dlq_durable = format!("integration-cdlq-{batch}");
    let dlq_handle = tokio::spawn(async move {
        dlq_bus
            .consume_envelopes::<DeadLetterPayload, _, _>(
                ConsumerSpec {
                    stream: dlq_stream,
                    durable_name: dlq_durable,
                    filter_subject: "dlq.>".to_string(),
                    ack_wait: dlq_ack_wait,
                    max_deliver: dlq_max_deliver,
                },
                move |message| {
                    let dlq_tx = dlq_tx.clone();
                    async move {
                        let _ = dlq_tx.send(message.payload.original_message_id).await;
                        Ok(())
                    }
                },
            )
            .await
    });

    // 失败 consumer（总是失败，触发 redelivery -> DLQ）
    let retry_bus = bus.clone();
    let retry_spec = ConsumerSpec {
        stream: names.name(StreamKind::Inbound),
        durable_name: format!("integration-cdlq-retry-{batch}"),
        filter_subject: format!("inbound.wechat.{account_id}"),
        ack_wait: config.ack_wait(),
        max_deliver: config.max_deliver,
    };
    let retry_handle = tokio::spawn(async move {
        retry_bus
            .consume_envelopes::<InboundPayload, _, _>(retry_spec, |_| async move {
                anyhow::bail!("force redelivery")
            })
            .await
    });

    // 收集 count 条 DLQ 消息，验证无碰撞丢失
    let mut received = Vec::new();
    for _ in 0..count {
        let id = tokio::time::timeout(std::time::Duration::from_secs(30), dlq_rx.recv())
            .await?
            .ok_or_else(|| anyhow::anyhow!("dlq consumer closed early"))?;
        if let Some(id) = id {
            received.push(id);
        }
    }
    retry_handle.abort();
    dlq_handle.abort();

    received.sort();
    message_ids.sort();
    assert_eq!(received.len(), count, "DLQ 丢失了消息");
    for id in &message_ids {
        assert!(received.contains(id), "DLQ 缺少消息 {id}");
    }
    Ok(())
}