Add tests for audit logs partitions

application/src/test/java/org/thingsboard/server/controller/BaseAuditLogControllerTest.java

@@ -20,18 +20,33 @@ import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.boot.test.mock.mockito.SpyBean;
 import org.thingsboard.server.common.data.Device;
 import org.thingsboard.server.common.data.Tenant;
 import org.thingsboard.server.common.data.User;
+import org.thingsboard.server.common.data.audit.ActionType;
 import org.thingsboard.server.common.data.audit.AuditLog;
+import org.thingsboard.server.common.data.id.EntityId;
 import org.thingsboard.server.common.data.page.PageData;
 import org.thingsboard.server.common.data.page.TimePageLink;
 import org.thingsboard.server.common.data.security.Authority;
+import org.thingsboard.server.dao.audit.AuditLogDao;
 import org.thingsboard.server.dao.model.ModelConstants;
+import org.thingsboard.server.dao.sqlts.insert.sql.SqlPartitioningRepository;
+import org.thingsboard.server.service.ttl.AuditLogsCleanUpService;
 
+import java.time.LocalDate;
+import java.time.ZoneOffset;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.verify;
 import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
 
 public abstract class BaseAuditLogControllerTest extends AbstractControllerTest {
@@ -39,6 +54,18 @@ public abstract class BaseAuditLogControllerTest extends AbstractControllerTest
     private Tenant savedTenant;
     private User tenantAdmin;
 
+    @Autowired
+    private AuditLogDao auditLogDao;
+    @SpyBean
+    private SqlPartitioningRepository partitioningRepository;
+    @Autowired
+    private AuditLogsCleanUpService auditLogsCleanUpService;
+
+    @Value("#{${sql.audit_logs.partition_size} * 60 * 60 * 1000}")
+    private long partitionDurationInMs;
+    @Value("${sql.ttl.audit_logs.ttl}")
+    private long auditLogsTtlInSec;
+
     @Before
     public void beforeTest() throws Exception {
         loginSysAdmin();
@@ -145,4 +172,45 @@ public abstract class BaseAuditLogControllerTest extends AbstractControllerTest
 
         Assert.assertEquals(179, loadedAuditLogs.size());
     }
+
+    @Test
+    public void whenSavingNewAuditLog_thenCheckAndCreatePartitionIfNotExists() {
+        reset(partitioningRepository);
+        AuditLog auditLog = createAuditLog(ActionType.LOGIN, tenantAdminUserId);
+        verify(partitioningRepository).createPartitionIfNotExists(eq("audit_log"), eq(auditLog.getCreatedTime()), eq(partitionDurationInMs));
+
+        List<Long> partitions = partitioningRepository.fetchPartitions("audit_log");
+        assertThat(partitions).singleElement().satisfies(partitionStartTs -> {
+            assertThat(partitionStartTs).isEqualTo(partitioningRepository.calculatePartitionStartTime(auditLog.getCreatedTime(), partitionDurationInMs));
+        });
+    }
+
+    @Test
+    public void whenCleaningUpAuditLogsByTtl_thenDropOldPartitions() {
+        long oldAuditLogTs = LocalDate.of(2020, 10, 1).atStartOfDay().toInstant(ZoneOffset.UTC).toEpochMilli();
+        long partitionStartTs = partitioningRepository.calculatePartitionStartTime(oldAuditLogTs, partitionDurationInMs);
+        partitioningRepository.createPartitionIfNotExists("audit_log", oldAuditLogTs, partitionDurationInMs);
+        List<Long> partitions = partitioningRepository.fetchPartitions("audit_log");
+        assertThat(partitions).contains(partitionStartTs);
+
+        auditLogsCleanUpService.cleanUp();
+        partitions = partitioningRepository.fetchPartitions("audit_log");
+        assertThat(partitions).doesNotContain(partitionStartTs);
+        assertThat(partitions).allSatisfy(partitionsStart -> {
+            long partitionEndTs = partitionsStart + partitionDurationInMs;
+            assertThat(partitionEndTs).isGreaterThan(System.currentTimeMillis() - TimeUnit.SECONDS.toMillis(auditLogsTtlInSec));
+        });
+    }
+
+    private AuditLog createAuditLog(ActionType actionType, EntityId entityId) {
+        AuditLog auditLog = new AuditLog();
+        auditLog.setTenantId(tenantId);
+        auditLog.setCustomerId(null);
+        auditLog.setUserId(tenantAdminUserId);
+        auditLog.setEntityId(entityId);
+        auditLog.setUserName(tenantAdmin.getEmail());
+        auditLog.setActionType(actionType);
+        return auditLogDao.save(tenantId, auditLog);
+    }
+
 }

application/src/test/resources/application-test.properties

@@ -56,4 +56,7 @@ queue.rule-engine.queues[2].processing-strategy.retries=1
 queue.rule-engine.queues[2].processing-strategy.pause-between-retries=0
 queue.rule-engine.queues[2].processing-strategy.max-pause-between-retries=0
 
-usage.stats.report.enabled=false
+usage.stats.report.enabled=false
+
+sql.audit_logs.partition_size=24
+sql.ttl.audit_logs.ttl=2592000

dao/src/main/java/org/thingsboard/server/dao/sqlts/insert/sql/SqlPartitioningRepository.java

@@ -127,7 +127,7 @@ public class SqlPartitioningRepository {
         return false;
     }
 
-    private List<Long> fetchPartitions(String table) {
+    public List<Long> fetchPartitions(String table) {
         List<Long> partitions = new ArrayList<>();
         List<String> partitionsTables = jdbcTemplate.queryForList(SELECT_PARTITIONS_STMT, new Object[]{table}, String.class);
         for (String partitionTableName : partitionsTables) {
@@ -141,7 +141,7 @@ public class SqlPartitioningRepository {
         return partitions;
     }
 
-    private long calculatePartitionStartTime(long ts, long partitionDuration) {
+    public long calculatePartitionStartTime(long ts, long partitionDuration) {
         return ts - (ts % partitionDuration);
     }