content: /**
   * Copyright 2020 Alibaba Group Holding Limited.
   *
   * <p>Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
   * except in compliance with the License. You may obtain a copy of the License at
   *
   * <p>http://www.apache.org/licenses/LICENSE-2.0
   *
   * <p>Unless required by applicable law or agreed to in writing, software distributed under the
   * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
   * express or implied. See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package com.alibaba.graphscope.groot.tests.coordinator;
  
  import static com.alibaba.graphscope.groot.coordinator.backup.BackupManager.*;
  
  import static org.junit.jupiter.api.Assertions.*;
  import static org.mockito.ArgumentMatchers.anyString;
  import static org.mockito.Mockito.*;
  
  import com.alibaba.graphscope.groot.CompletionCallback;
  import com.alibaba.graphscope.groot.SnapshotWithSchema;
  import com.alibaba.graphscope.groot.common.config.BackupConfig;
  import com.alibaba.graphscope.groot.common.config.CommonConfig;
  import com.alibaba.graphscope.groot.common.config.Configs;
  import com.alibaba.graphscope.groot.common.schema.wrapper.GraphDef;
  import com.alibaba.graphscope.groot.common.util.BackupInfo;
  import com.alibaba.graphscope.groot.coordinator.*;
  import com.alibaba.graphscope.groot.coordinator.backup.BackupManager;
  import com.alibaba.graphscope.groot.coordinator.backup.StoreBackupTaskSender;
  import com.alibaba.graphscope.groot.frontend.SnapshotCache;
  import com.alibaba.graphscope.groot.meta.MetaService;
  import com.alibaba.graphscope.groot.meta.MetaStore;
  import com.alibaba.graphscope.groot.store.backup.StoreBackupId;
  import com.fasterxml.jackson.core.type.TypeReference;
  import com.fasterxml.jackson.databind.ObjectMapper;
  
  import org.apache.commons.io.FileUtils;
  import org.junit.jupiter.api.Test;
  
  import java.io.File;
  import java.io.IOException;
  import java.util.*;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.TimeUnit;
  
  public class BackupManagerTest {
  
      @Test
      void testBackupManager() throws IOException, InterruptedException {
          // init config
          Configs configs =
                  Configs.newBuilder()
                          .put(CommonConfig.STORE_NODE_COUNT.getKey(), "2")
                          .put(BackupConfig.BACKUP_ENABLE.getKey(), "true")
                          .put(BackupConfig.BACKUP_CREATION_BUFFER_MAX_COUNT.getKey(), "4")
                          .put(BackupConfig.BACKUP_GC_INTERVAL_HOURS.getKey(), "24")
                          .put(BackupConfig.BACKUP_AUTO_SUBMIT.getKey(), "false")
                          .put(BackupConfig.BACKUP_AUTO_SUBMIT_INTERVAL_HOURS.getKey(), "24")
                          .build();
  
          // init data
          long querySnapshotId = 10L;
          GraphDef graphDef = GraphDef.newBuilder().setVersion(10L).build();
          SnapshotWithSchema snapshotWithSchema = new SnapshotWithSchema(querySnapshotId, graphDef);
          List<Long> queueOffsets = new ArrayList<>();
          Map<Integer, Integer> partitionToBackupId1 = new HashMap<>();
          partitionToBackupId1.put(0, 1);
          partitionToBackupId1.put(1, 1);
          Map<Integer, Integer> partitionToBackupId2 = new HashMap<>();
          partitionToBackupId2.put(0, 2);
          partitionToBackupId2.put(1, 2);
          BackupInfo backupInfo1 =
                  new BackupInfo(
                          1,
                          querySnapshotId,
                          graphDef.toProto().toByteArray(),
                          queueOffsets,
                          partitionToBackupId1);
          BackupInfo backupInfo2 =
                  new BackupInfo(
                          2,
                          querySnapshotId,
                          graphDef.toProto().toByteArray(),
                          queueOffsets,
                          partitionToBackupId2);
  
          // mock MetaStore behaviours
          ObjectMapper objectMapper = new ObjectMapper();
          MetaStore mockMetaStore = mock(MetaStore.class);
          when(mockMetaStore.exists(anyString())).thenReturn(true);
          when(mockMetaStore.read(GLOBAL_BACKUP_ID_PATH))
                  .thenReturn(objectMapper.writeValueAsBytes(0));
          when(mockMetaStore.read(BACKUP_INFO_PATH))
                  .thenReturn(objectMapper.writeValueAsBytes(new ArrayList<BackupInfo>()));
  
          // mock MetaService behaviours
          MetaService mockMetaService = mock(MetaService.class);
          when(mockMetaService.getPartitionCount()).thenReturn(2);
          when(mockMetaService.getStoreIdByPartition(0)).thenReturn(0);
          when(mockMetaService.getStoreIdByPartition(1)).thenReturn(1);
  
          // mock SnapshotManager behaviours
          SnapshotManager mockSnapshotManager = mock(SnapshotManager.class);
          when(mockSnapshotManager.getQueueOffsets()).thenReturn(queueOffsets);
  
          // mock SchemaManager
          SchemaManager mockSchemaManager = mock(SchemaManager.class);
          when(mockSchemaManager.getGraphDef()).thenReturn(graphDef);
  
          // mock SnapshotCache
          SnapshotCache mockSnapshotCache = mock(SnapshotCache.class);
          when(mockSnapshotCache.getSnapshotWithSchema()).thenReturn(snapshotWithSchema);
  
          // mock StoreBackupTaskSender behaviours
          StoreBackupTaskSender mockStoreBackupTaskSender = mock(StoreBackupTaskSender.class);
          doAnswer(
                          invocation -> {
                              int partitionOrStoreId = invocation.getArgument(0);
                              int globalBackupId = invocation.getArgument(1);
                              CompletionCallback<StoreBackupId> callback = invocation.getArgument(2);
                              StoreBackupId storeBackupId = new StoreBackupId(globalBackupId);
                              storeBackupId.addPartitionBackupId(partitionOrStoreId, globalBackupId);
                              callback.onCompleted(storeBackupId);
                              return null;
                          })
                  .when(mockStoreBackupTaskSender)
                  .createStoreBackup(anyInt(), anyInt(), any());
          doAnswer(
                          invocation -> {
                              CompletionCallback<Void> callback = invocation.getArgument(3);
                              callback.onCompleted(null);
                              return null;
                          })
                  .when(mockStoreBackupTaskSender)
                  .restoreFromStoreBackup(anyInt(), any(), anyString(), any());
          doAnswer(
                          invocation -> {
                              CompletionCallback<Void> callback = invocation.getArgument(2);
                              callback.onCompleted(null);
                              return null;
                          })
                  .when(mockStoreBackupTaskSender)
                  .verifyStoreBackup(anyInt(), any(), any());
  
          BackupManager backupManager =
                  new BackupManager(
                          configs,
                          mockMetaService,
                          mockMetaStore,
                          mockSnapshotManager,
                          mockSchemaManager,
                          mockSnapshotCache,
                          mockStoreBackupTaskSender);
          backupManager.start();
          verify(mockSnapshotManager).addListener(any());
  
          // create the first backup
          CountDownLatch updateBackupIdLatch1 = new CountDownLatch(1);
          CountDownLatch updateBackupInfoByCreation1Latch = new CountDownLatch(1);
          doAnswer(
                          invocation -> {
                              updateBackupIdLatch1.countDown();
                              return null;
                          })
                  .when(mockMetaStore)
                  .write(GLOBAL_BACKUP_ID_PATH, objectMapper.writeValueAsBytes(1));
          doAnswer(
                          invocation -> {
                              byte[] backupInfoBytes = invocation.getArgument(1);
                              List<BackupInfo> backupInfoList =
                                      objectMapper.readValue(
                                              backupInfoBytes,
                                              new TypeReference<List<BackupInfo>>() {});
                              assertEquals(backupInfoList.size(), 1);
                              assertEquals(backupInfoList.get(0), backupInfo1);
                              updateBackupInfoByCreation1Latch.countDown();
                              return null;
                          })
                  .when(mockMetaStore)
                  .write(
                          BACKUP_INFO_PATH,
                          objectMapper.writeValueAsBytes(Collections.singletonList(backupInfo1)));
          int backupId1 = backupManager.createNewBackup();
          assertEquals(backupId1, 1);
          assertTrue(updateBackupIdLatch1.await(5L, TimeUnit.SECONDS));
          assertTrue(updateBackupInfoByCreation1Latch.await(5L, TimeUnit.SECONDS));
  
          // create the second backup
          CountDownLatch updateBackupIdLatch2 = new CountDownLatch(1);
          CountDownLatch updateBackupInfoByCreation2Latch = new CountDownLatch(1);
          doAnswer(
                          invocation -> {
                              updateBackupIdLatch2.countDown();
                              return null;
                          })
                  .when(mockMetaStore)
                  .write(GLOBAL_BACKUP_ID_PATH, objectMapper.writeValueAsBytes(2));
          doAnswer(
                          invocation -> {
                              byte[] backupInfoBytes = invocation.getArgument(1);
                              List<BackupInfo> backupInfoList =
                                      objectMapper.readValue(
                                              backupInfoBytes,
                                              new TypeReference<List<BackupInfo>>() {});
                              backupInfoList.sort(
                                      new Comparator<BackupInfo>() {
                                          @Override
                                          public int compare(BackupInfo o1, BackupInfo o2) {
                                              return o1.getGlobalBackupId() - o2.getGlobalBackupId();
                                          }
                                      });
                              assertEquals(backupInfoList.size(), 2);
                              assertEquals(backupInfoList.get(0), backupInfo1);
                              assertEquals(backupInfoList.get(1), backupInfo2);
                              updateBackupInfoByCreation2Latch.countDown();
                              return null;
                          })
                  .when(mockMetaStore)
                  .write(
                          BACKUP_INFO_PATH,
                          objectMapper.writeValueAsBytes(Arrays.asList(backupInfo1, backupInfo2)));
          int backupId2 = backupManager.createNewBackup();
          assertEquals(backupId2, 2);
          assertTrue(updateBackupIdLatch2.await(5L, TimeUnit.SECONDS));
          assertTrue(updateBackupInfoByCreation2Latch.await(5L, TimeUnit.SECONDS));
  
          // get backup info list and check
          assertEquals(backupManager.getBackupInfoList().size(), 2);
  
          // verify backups
          try {
              backupManager.verifyBackup(backupId1);
              backupManager.verifyBackup(backupId2);
          } catch (Exception e) {
              fail("should not have thrown any exception during backup verification");
          }
  
          // restore from the second backup
          try {
              // TODO(siyuan): Fix this test
              backupManager.restoreFromBackup(backupId2, "restore_meta", "restore_store");
              // assertTrue(Files.exists(Paths.get("restore_meta", "query_snapshot_id")));
              // assertTrue(Files.exists(Paths.get("restore_meta", "graph_def_proto_bytes")));
              // assertTrue(Files.exists(Paths.get("restore_meta", "queue_offsets")));
          } catch (Exception e) {
              fail("should not have thrown any exception during backup restoring");
          } finally {
              FileUtils.deleteDirectory(new File("restore_meta"));
          }
  
          // purge 1 old backup
          CountDownLatch updateBackupInfoByPurgingLatch = new CountDownLatch(1);
          doAnswer(
                          invocation -> {
                              byte[] backupInfoBytes = invocation.getArgument(1);
                              List<BackupInfo> backupInfoList =
                                      objectMapper.readValue(
                                              backupInfoBytes,
                                              new TypeReference<List<BackupInfo>>() {});
                              assertEquals(backupInfoList.size(), 1);
                              assertEquals(backupInfoList.get(0), backupInfo2);
                              updateBackupInfoByPurgingLatch.countDown();
                              return null;
                          })
                  .when(mockMetaStore)
                  .write(
                          BACKUP_INFO_PATH,
                          objectMapper.writeValueAsBytes(Collections.singletonList(backupInfo2)));
          backupManager.purgeOldBackups(1);
          assertTrue(updateBackupInfoByPurgingLatch.await(5L, TimeUnit.SECONDS));
  
          // get backup info list and check
          assertEquals(backupManager.getBackupInfoList().size(), 1);
  
          // delete the remaining backup '2'
          CountDownLatch updateBackupInfoByDeletionLatch = new CountDownLatch(1);
          doAnswer(
                          invocation -> {
                              byte[] backupInfoBytes = invocation.getArgument(1);
                              List<BackupInfo> backupInfoList =
                                      objectMapper.readValue(
                                              backupInfoBytes,
                                              new TypeReference<List<BackupInfo>>() {});
                              assertTrue(backupInfoList.isEmpty());
                              updateBackupInfoByDeletionLatch.countDown();
                              return null;
                          })
                  .when(mockMetaStore)
                  .write(
                          BACKUP_INFO_PATH,
                          objectMapper.writeValueAsBytes(new ArrayList<BackupInfo>()));
          backupManager.deleteBackup(2);
          assertTrue(updateBackupInfoByDeletionLatch.await(5L, TimeUnit.SECONDS));
  
          // get backup info list and check
          assertTrue(backupManager.getBackupInfoList().isEmpty());
  
          backupManager.stop();
          verify(mockSnapshotManager).removeListener(any());
      }
  }
  