content: /**
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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 org.apache.atlas.repository.store.graph.v2;
  
  import com.fasterxml.jackson.core.type.TypeReference;
  import org.apache.atlas.exception.AtlasBaseException;
  import org.apache.atlas.model.instance.AtlasEntityHeader;
  import org.apache.atlas.model.instance.AtlasEntityHeaders;
  import org.apache.atlas.model.typedef.AtlasEntityDef;
  import org.apache.atlas.repository.audit.EntityAuditRepository;
  import org.apache.atlas.repository.graphdb.AtlasGraph;
  import org.apache.atlas.repository.store.graph.AtlasEntityStore;
  import org.apache.atlas.type.AtlasEntityType;
  import org.apache.atlas.type.AtlasTypeRegistry;
  import org.apache.atlas.utils.AtlasJson;
  import org.apache.atlas.utils.TestResourceFileUtils;
  import org.apache.commons.lang.StringUtils;
  import org.elasticsearch.common.util.CollectionUtils;
  import org.testng.annotations.Test;
  
  import java.io.IOException;
  import java.util.HashSet;
  import java.util.Map;
  import java.util.Set;
  
  import static org.apache.atlas.repository.store.graph.v2.ClassificationAssociator.Updater.PROCESS_ADD;
  import static org.apache.atlas.repository.store.graph.v2.ClassificationAssociator.Updater.PROCESS_DELETE;
  import static org.apache.atlas.repository.store.graph.v2.ClassificationAssociator.Updater.PROCESS_UPDATE;
  import static org.apache.atlas.repository.store.graph.v2.ClassificationAssociator.Updater.STATUS_DONE;
  import static org.apache.atlas.repository.store.graph.v2.ClassificationAssociator.Updater.STATUS_SKIPPED;
  import static org.mockito.Matchers.anyLong;
  import static org.mockito.Matchers.anyString;
  import static org.mockito.Mockito.mock;
  import static org.mockito.Mockito.when;
  import static org.testng.Assert.assertEquals;
  import static org.testng.Assert.assertNotNull;
  import static org.testng.Assert.assertTrue;
  import static org.testng.FileAssert.fail;
  
  public class ClassificationAssociatorTest {
      private static final String TABLE_GUID            = "df122fc3-5555-40f8-a30f-3090b8a622f8";
      private static final String TEST_FILES_SUBDIR     = "classification-association";
      private static final String MESSAGE_SEPARATOR     = ":";
      private static final String ENTITY_NAME_SEPARATOR = "->";
  
      @Test
      public void auditScanYieldsNothing_EmptyHeadersReturned() {
          AtlasEntityHeaders actualEntityHeaders = setupRetriever("header-empty", 0, 0, null);
  
          assertNotNull(actualEntityHeaders);
          assertEquals(actualEntityHeaders.getGuidHeaderMap().size(), 0);
      }
  
      @Test
      public void auditScanYieldsOneEntity_EntityHeadersHasOneElementWithClassification() {
          AtlasEntityHeaders actualEntityHeaders = setupRetriever("header-Tx", 0, 0, TABLE_GUID);
  
          assertNotNull(actualEntityHeaders);
          assertEquals(actualEntityHeaders.getGuidHeaderMap().size(), 1);
          assertTrue(actualEntityHeaders.getGuidHeaderMap().containsKey(TABLE_GUID));
          assertEquals(actualEntityHeaders.getGuidHeaderMap().get(TABLE_GUID).getGuid(), TABLE_GUID);
          assertNotNull(actualEntityHeaders.getGuidHeaderMap().get(TABLE_GUID).getClassifications());
          assertEquals(actualEntityHeaders.getGuidHeaderMap().get(TABLE_GUID).getClassifications().size(), 1);
      }
  
      @Test
      public void updaterIncorrectType_ReturnsError() throws IOException {
          AtlasEntityHeaders entityHeaderMap = getEntityHeaderMapFromFile("header-PII");
          AtlasEntityStore   entitiesStore   = mock(AtlasEntityStore.class);
          AtlasTypeRegistry  typeRegistry    = mock(AtlasTypeRegistry.class);
  
          when(typeRegistry.getEntityTypeByName(anyString())).thenReturn(null);
  
          AtlasGraph                       atlasGraph = mock(AtlasGraph.class);
          ClassificationAssociator.Updater updater    = new ClassificationAssociator.Updater(atlasGraph, typeRegistry, entitiesStore);
          String                           summary    = updater.setClassifications(entityHeaderMap.getGuidHeaderMap());
  
          assertTrue(summary.contains("hive_"));
          assertTrue(summary.contains(STATUS_SKIPPED));
      }
  
      @Test
      public void updaterCorrectTypeEntityNotFound_Skipped() throws IOException {
          AtlasEntityHeaders entityHeaderMap = getEntityHeaderMapFromFile("header-PII");
          AtlasEntityType    hiveTable       = mock(AtlasEntityType.class);
          AtlasEntityStore   entitiesStore   = mock(AtlasEntityStore.class);
          AtlasTypeRegistry  typeRegistry    = mock(AtlasTypeRegistry.class);
          AtlasGraph         atlasGraph      = mock(AtlasGraph.class);
  
          when(typeRegistry.getEntityTypeByName(anyString())).thenReturn(hiveTable);
          when(hiveTable.getTypeName()).thenReturn("hive_column");
  
          ClassificationAssociatorUpdaterForSpy updater = new ClassificationAssociatorUpdaterForSpy(atlasGraph, typeRegistry, entitiesStore);
          String                                summary = updater.setClassifications(entityHeaderMap.getGuidHeaderMap());
  
          TypeReference<String[]> typeReference = new TypeReference<String[]>() {};
          String[]                summaryArray  = AtlasJson.fromJson(summary, typeReference);
  
          assertEquals(summaryArray.length, 1);
          assertSummaryElement(summaryArray[0], "Entity", STATUS_SKIPPED, "");
      }
  
      @Test
      public void updaterEntityWithUniqueName() throws IOException, AtlasBaseException {
          AtlasEntityDef ed = getAtlasEntityDefFromFile("col-entity-def-unique-name");
  
          AtlasEntityHeaders                           entityHeaderMap = getEntityHeaderMapFromFile("header-PII-no-qualifiedName");
          AtlasEntityStore                             entitiesStore   = mock(AtlasEntityStore.class);
          AtlasTypeRegistry                            typeRegistry    = new AtlasTypeRegistry();
          AtlasTypeRegistry.AtlasTransientTypeRegistry ttr             = typeRegistry.lockTypeRegistryForUpdate();
  
          ttr.addTypes(CollectionUtils.newSingletonArrayList(ed));
  
          AtlasGraph                            atlasGraph = mock(AtlasGraph.class);
          ClassificationAssociatorUpdaterForSpy updater    = new ClassificationAssociatorUpdaterForSpy(atlasGraph, ttr, entitiesStore, "col-entity-PII");
          String                                summary    = updater.setClassifications(entityHeaderMap.getGuidHeaderMap());
  
          TypeReference<String[]> typeReference = new TypeReference<String[]>() {};
          String[]                summaryArray  = AtlasJson.fromJson(summary, typeReference);
  
          assertEquals(summaryArray.length, 1);
          assertSummaryElement(summaryArray[0], "Update", STATUS_DONE, "PII");
      }
  
      @Test
      public void updaterTests() throws IOException {
          updaterAssert("header-None", "col-entity-None");
          updaterAssert("header-PII", "col-entity-None", PROCESS_ADD + ":PII");
          updaterAssert("header-PII", "col-entity-PII", PROCESS_UPDATE + ":PII");
          updaterAssert("header-None", "col-entity-PII", PROCESS_DELETE + ":PII");
          updaterAssert("header-PII-VENDOR_PII", "col-entity-PII-FIN_PII", PROCESS_DELETE + ":FIN_PII", PROCESS_UPDATE + ":PII", PROCESS_ADD + ":VENDOR_PII");
          updaterAssert("header-None", "col-entity-None");
          updaterAssert("header-FIN_PII", "col-entity-PII", PROCESS_DELETE + ":PII", PROCESS_ADD + ":FIN_PII");
      }
  
      @Test
      public void updater_filterPropagatedClassifications() throws IOException {
          updaterAssert("header-Tx-prop-T1", "col-entity-T1-prop-Tn", PROCESS_DELETE + ":T1", PROCESS_ADD + ":Tx");
          updaterAssert("header-Tx-prop-T1-No-Guid", "col-entity-T1-prop-Tn-No-Guid", PROCESS_DELETE + ":Tn", PROCESS_UPDATE + ":T1", PROCESS_ADD + ":Tx");
      }
  
      private AtlasEntityHeaders setupRetriever(String headersFile, int fromTimestamp, int toTimestamp, final String tableGuid) {
          AtlasEntityHeader entityHeaderWithClassification;
  
          try {
              Set<String> guids = new HashSet<>();
  
              entityHeaderWithClassification = TestResourceFileUtils.readObjectFromJson(TEST_FILES_SUBDIR, headersFile, AtlasEntityHeader.class);
  
              if (!StringUtils.isEmpty(tableGuid)) {
                  guids.add(tableGuid);
              }
  
              EntityAuditRepository auditRepository = mock(EntityAuditRepository.class);
  
              when(auditRepository.getEntitiesWithTagChanges(anyLong(), anyLong())).thenReturn(guids);
  
              EntityGraphRetriever entityGraphRetriever = mock(EntityGraphRetriever.class);
  
              when(entityGraphRetriever.toAtlasEntityHeaderWithClassifications(TABLE_GUID)).thenReturn(entityHeaderWithClassification);
  
              ClassificationAssociator.Retriever retriever = new ClassificationAssociator.Retriever(entityGraphRetriever, auditRepository);
  
              return retriever.get(fromTimestamp, toTimestamp);
          } catch (Exception ex) {
              fail("Exception!");
              return null;
          }
      }
  
      private void assertSummaryElement(String summaryElement, String operation, String status, String classificationName) {
          String[] splits     = StringUtils.split(summaryElement, MESSAGE_SEPARATOR);
          String[] nameSplits = StringUtils.split(splits[3], ENTITY_NAME_SEPARATOR);
  
          if (nameSplits.length > 1) {
              assertEquals(nameSplits[1].trim(), classificationName);
          }
  
          assertEquals(splits[0], operation);
          assertEquals(splits[4], status);
      }
  
      private String[] setupUpdater(String entityHeaderFileName, String entityFileName, int expectedSummaryLength) throws IOException {
          AtlasEntityHeaders entityHeaderMap = getEntityHeaderMapFromFile(entityHeaderFileName);
          AtlasEntityType    hiveTable       = mock(AtlasEntityType.class);
          AtlasEntityStore   entitiesStore   = mock(AtlasEntityStore.class);
          AtlasTypeRegistry  typeRegistry    = mock(AtlasTypeRegistry.class);
          AtlasGraph         atlasGraph      = mock(AtlasGraph.class);
  
          when(typeRegistry.getEntityTypeByName(anyString())).thenReturn(hiveTable);
          when(hiveTable.getTypeName()).thenReturn("hive_column");
  
          ClassificationAssociatorUpdaterForSpy updater = new ClassificationAssociatorUpdaterForSpy(atlasGraph, typeRegistry, entitiesStore, entityFileName);
          String                                summary = updater.setClassifications(entityHeaderMap.getGuidHeaderMap());
  
          TypeReference<String[]> typeReference = new TypeReference<String[]>() {};
          String[]                summaryArray  = AtlasJson.fromJson(summary, typeReference);
  
          assertEquals(summaryArray.length, expectedSummaryLength);
  
          return summaryArray;
      }
  
      private AtlasEntityHeader getEntityHeaderFromFile(String entityJson) throws IOException {
          return TestResourceFileUtils.readObjectFromJson(TEST_FILES_SUBDIR, entityJson, AtlasEntityHeader.class);
      }
  
      private AtlasEntityHeaders getEntityHeaderMapFromFile(String filename) throws IOException {
          return TestResourceFileUtils.readObjectFromJson(TEST_FILES_SUBDIR, filename, AtlasEntityHeaders.class);
      }
  
      private AtlasEntityDef getAtlasEntityDefFromFile(String filename) throws IOException {
          return TestResourceFileUtils.readObjectFromJson(TEST_FILES_SUBDIR, filename, AtlasEntityDef.class);
      }
  
      private void updaterAssert(String incoming, String entity, String... opNamePair) throws IOException {
          String[] summary = setupUpdater(incoming, entity, opNamePair.length);
  
          for (int i = 0; i < opNamePair.length; i++) {
              String   s      = opNamePair[i];
              String[] splits = StringUtils.split(s, ":");
  
              assertSummaryElement(summary[i], splits[0], STATUS_DONE, splits[1]);
          }
      }
  
      private class ClassificationAssociatorUpdaterForSpy extends ClassificationAssociator.Updater {
          private final String entityFileName;
  
          public ClassificationAssociatorUpdaterForSpy(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore) {
              super(atlasGraph, typeRegistry, entitiesStore);
              this.entityFileName = StringUtils.EMPTY;
          }
  
          public ClassificationAssociatorUpdaterForSpy(AtlasGraph atlasGraph, AtlasTypeRegistry typeRegistry, AtlasEntityStore entitiesStore, String entityFileName) {
              super(atlasGraph, typeRegistry, entitiesStore);
  
              this.entityFileName = entityFileName;
          }
  
          @Override
          AtlasEntityHeader getByUniqueAttributes(AtlasEntityType entityType, String qualifiedName, Map<String, Object> attrValues) {
              try {
                  if (StringUtils.isEmpty(entityFileName)) {
                      return null;
                  }
  
                  return getEntityHeaderFromFile(entityFileName);
              } catch (IOException e) {
                  fail(entityFileName + " could not be loaded.");
                  return null;
              }
          }
      }
  }
  