ai-security-xdr/haobang-security-dm/syslog-consumer-rule/src/test/java/com/common/service/BacklogSamplingAuditServiceTest.java

package com.common.service;

import com.common.mapper.BacklogRetentionAuditMapper;
import com.config.BacklogSamplingProperties;
import org.junit.jupiter.api.Test;
import org.springframework.test.util.ReflectionTestUtils;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.TransactionDefinition;
import org.springframework.transaction.TransactionStatus;

import java.time.LocalDateTime;
import java.util.Arrays;
import java.util.List;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

class BacklogSamplingAuditServiceTest {

    @Test
    void aggregatesSameHourTopicPartitionAndPolicyWithoutMutatingInput() {
        BacklogRetentionDelta first = delta(1, 1, 0, 1, 0);
        BacklogRetentionDelta second = delta(1, 0, 1, 0, 1);

        List<BacklogRetentionDelta> result = BacklogSamplingAuditService.aggregate(
                Arrays.asList(first, second));

        assertEquals(1, result.size());
        assertEquals(2L, result.get(0).getMessageTotal());
        assertEquals(1L, result.get(0).getNormalRowsSaved());
        assertEquals(1L, result.get(0).getNormalRowsDropped());
        assertEquals(1L, result.get(0).getRawSaved());
        assertEquals(1L, result.get(0).getRawDropped());
        assertEquals(1L, first.getMessageTotal(), "聚合不能修改调用方的逐消息审计对象");
    }

    @Test
    void auditFailureReturnsFalseAndRecordsFullRetentionFallback() {
        BacklogRetentionAuditMapper mapper = mock(BacklogRetentionAuditMapper.class);
        when(mapper.upsert(any())).thenThrow(new IllegalStateException("table unavailable"));
        PlatformTransactionManager transactionManager = transactionManager();
        BacklogSamplingMetrics metrics = new BacklogSamplingMetrics();
        BacklogSamplingAuditService service = new BacklogSamplingAuditService(
                mapper, properties(), metrics, transactionManager);

        boolean audited = service.recordBeforeDropping(Arrays.asList(
                delta(2, 1, 1, 1, 1)));

        assertFalse(audited);
        Map<String, Object> snapshot = metrics.snapshot();
        assertEquals(1L, snapshot.get("audit_write_failures_total"));
        assertEquals(2L, snapshot.get("audit_fallback_messages_total"));
    }

    @Test
    void successfulAuditCommitsAndRecordsSamplingMetrics() {
        BacklogRetentionAuditMapper mapper = mock(BacklogRetentionAuditMapper.class);
        when(mapper.upsert(any())).thenReturn(1);
        PlatformTransactionManager transactionManager = transactionManager();
        BacklogSamplingMetrics metrics = new BacklogSamplingMetrics();
        BacklogSamplingAuditService service = new BacklogSamplingAuditService(
                mapper, properties(), metrics, transactionManager);

        assertTrue(service.recordBeforeDropping(Arrays.asList(
                delta(1, 0, 1, 0, 1))));

        verify(mapper).upsert(any());
        assertEquals(1L, metrics.snapshot().get("degraded_messages_total"));
        assertEquals(1L, metrics.snapshot().get("normal_rows_dropped_total"));
    }

    private BacklogSamplingProperties properties() {
        BacklogSamplingProperties properties = new BacklogSamplingProperties();
        ReflectionTestUtils.setField(properties, "auditTimeoutSeconds", 5);
        ReflectionTestUtils.setField(properties, "auditRetryBackoffMs", 60_000L);
        return properties;
    }

    private PlatformTransactionManager transactionManager() {
        PlatformTransactionManager transactionManager = mock(PlatformTransactionManager.class);
        TransactionStatus status = mock(TransactionStatus.class);
        when(transactionManager.getTransaction(any(TransactionDefinition.class))).thenReturn(status);
        return transactionManager;
    }

    private BacklogRetentionDelta delta(long messages, long normalSaved, long normalDropped,
                                        long rawSaved, long rawDropped) {
        return new BacklogRetentionDelta(LocalDateTime.of(2026, 9, 12, 18, 0),
                "raw", 0, "v1", 1000, 10000,
                messages, 0, 0, 1, normalSaved, normalDropped,
                0, 0, 0, 1, rawSaved, rawDropped, 3_600_000L, 3_600_000L);
    }
}