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.migration;
  
  import com.google.inject.Inject;
  import org.apache.atlas.exception.AtlasBaseException;
  import org.apache.atlas.repository.AtlasTestBase;
  import org.apache.atlas.repository.graph.AtlasGraphProvider;
  import org.apache.atlas.repository.graph.GraphBackedSearchIndexer;
  import org.apache.atlas.repository.graph.GraphHelper;
  import org.apache.atlas.repository.graphdb.AtlasEdge;
  import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
  import org.apache.atlas.repository.graphdb.AtlasGraph;
  import org.apache.atlas.repository.graphdb.AtlasGraphQuery;
  import org.apache.atlas.repository.graphdb.AtlasVertex;
  import org.apache.atlas.repository.graphdb.GraphDBMigrator;
  import org.apache.atlas.repository.store.bootstrap.AtlasTypeDefStoreInitializer;
  import org.apache.atlas.repository.store.graph.v2.AtlasGraphUtilsV2;
  import org.apache.atlas.store.AtlasTypeDefStore;
  import org.apache.atlas.type.AtlasTypeRegistry;
  import org.apache.atlas.utils.TestResourceFileUtils;
  import org.apache.commons.lang.StringUtils;
  import org.testng.annotations.AfterClass;
  import org.testng.annotations.BeforeClass;
  
  import java.io.IOException;
  import java.util.Iterator;
  
  import static org.apache.atlas.utils.TestLoadModelUtils.loadModelFromJson;
  import static org.testng.Assert.assertEquals;
  import static org.testng.Assert.assertNotNull;
  import static org.testng.Assert.assertTrue;
  
  public class MigrationBaseAsserts extends AtlasTestBase {
      protected static final String ASSERT_NAME_PROPERTY = "Asset.name";
      private static final   String TYPE_NAME_PROPERTY   = "__typeName";
      private static final   String R_GUID_PROPERTY_NAME = "_r__guid";
      private final GraphDBMigrator migrator;
      private final AtlasGraph      graph;
  
      @Inject
      protected AtlasTypeDefStore typeDefStore;
  
      @Inject
      protected AtlasTypeRegistry typeRegistry;
  
      @Inject
      private AtlasTypeDefStoreInitializer storeInitializer;
  
      @Inject
      private GraphBackedSearchIndexer indexer;
  
      protected MigrationBaseAsserts(AtlasGraph graph, GraphDBMigrator migrator) {
          this.graph    = graph;
          this.migrator = migrator;
      }
  
      @BeforeClass
      public void initialize() throws Exception {
          super.initialize();
      }
  
      @AfterClass
      public void clear() throws Exception {
          AtlasGraphProvider.cleanup();
  
          super.cleanup();
      }
  
      protected void loadTypesFromJson() throws IOException, AtlasBaseException {
          loadModelFromJson("0000-Area0/0010-base_model.json", typeDefStore, typeRegistry);
          loadModelFromJson("1000-Hadoop/1020-fs_model.json", typeDefStore, typeRegistry);
          loadModelFromJson("1000-Hadoop/1030-hive_model.json", typeDefStore, typeRegistry);
      }
  
      protected void runFileImporter(String directoryToImport) throws IOException, AtlasBaseException {
          loadTypesFromJson();
          String directoryName = TestResourceFileUtils.getDirectory(directoryToImport);
          DataMigrationService.FileImporter fi = new DataMigrationService.FileImporter(migrator, typeDefStore, typeRegistry,
                  storeInitializer, directoryName, indexer);
  
          fi.run();
      }
  
      protected void assertHiveVertices(int dbCount, int tableCount, int columnCount) {
          int i = 0;
  
          Iterator<AtlasVertex> results = getVertices("hive_db", null);
          for (Iterator<AtlasVertex> it = results; it.hasNext(); i++) {
              assertNotNull(it.next());
          }
          assertEquals(i, dbCount);
  
          i       = 0;
          results = getVertices("hive_table", null);
          for (Iterator<AtlasVertex> it = results; it.hasNext(); i++) {
              assertNotNull(it.next());
          }
          assertEquals(i, tableCount);
  
          i       = 0;
          results = getVertices("hive_column", null);
          for (Iterator<AtlasVertex> it = results; it.hasNext(); i++) {
              assertNotNull(it.next());
          }
  
          assertTrue(i > 0);
          assertEquals(i, columnCount);
      }
  
      protected Iterator<AtlasVertex> getVertices(String typeName, String name) {
          AtlasGraphQuery query = graph.query().has(TYPE_NAME_PROPERTY, typeName);
  
          if (!StringUtils.isEmpty(name)) {
              query = query.has(ASSERT_NAME_PROPERTY, name);
          }
  
          return query.vertices().iterator();
      }
  
      protected AtlasVertex getVertex(String typeName, String name) {
          Iterator<AtlasVertex> iterator = getVertices(typeName, name);
  
          return iterator.hasNext() ? iterator.next() : null;
      }
  
      protected void assertEdges(String typeName, String assetName, AtlasEdgeDirection edgeDirection, int expectedItems, String edgeTypeName) {
          Iterator edgeIterator = getVertex(typeName, assetName).getEdges(edgeDirection).iterator();
          assertEdges(edgeIterator, expectedItems, edgeTypeName);
      }
  
      protected void assertEdges(Iterator<AtlasEdge> results, int expectedItems, String edgeTypeNameExpected) {
          int       count             = 0;
          AtlasEdge e                 = null;
          boolean   searchedEdgeFound = false;
          for (Iterator<AtlasEdge> it = results; it.hasNext(); ) {
              e = it.next();
              String typeName = AtlasGraphUtilsV2.getEncodedProperty(e, TYPE_NAME_PROPERTY, String.class);
              searchedEdgeFound = StringUtils.isEmpty(edgeTypeNameExpected) || typeName.equals(edgeTypeNameExpected);
              if (searchedEdgeFound) {
                  count++;
                  break;
              }
          }
  
          assertNotNull(AtlasGraphUtilsV2.getEncodedProperty(e, R_GUID_PROPERTY_NAME, Object.class));
          assertNotNull(AtlasGraphUtilsV2.getEncodedProperty(e, "tagPropagation", Object.class));
  
          if (StringUtils.isNotEmpty(edgeTypeNameExpected)) {
              assertTrue(searchedEdgeFound, edgeTypeNameExpected);
          }
  
          assertEquals(count, expectedItems, String.format("%s", edgeTypeNameExpected));
      }
  
      protected void assertEdgesWithLabel(Iterator<AtlasEdge> results, int startIdx, String edgeTypeName) {
          int       count = 0;
          AtlasEdge e     = null;
          for (Iterator<AtlasEdge> it = results; it.hasNext() && count < startIdx; count++) {
              e = it.next();
          }
  
          assertNotNull(AtlasGraphUtilsV2.getEncodedProperty(e, R_GUID_PROPERTY_NAME, Object.class));
          assertNotNull(AtlasGraphUtilsV2.getEncodedProperty(e, "tagPropagation", Object.class));
  
          if (StringUtils.isNotEmpty(edgeTypeName)) {
              assertEquals(e.getLabel(), edgeTypeName, edgeTypeName);
          }
      }
  
      protected void assertTypeCountNameGuid(String typeName, int expectedItems, String name, String guid) {
          Iterator<AtlasVertex> results = getVertices(typeName, name);
  
          int count = 0;
          for (Iterator<AtlasVertex> it = results; it.hasNext(); ) {
              AtlasVertex v = it.next();
  
              assertEquals(GraphHelper.getTypeName(v), typeName);
  
              if (StringUtils.isNotEmpty(guid)) {
                  assertEquals(GraphHelper.getGuid(v), guid, name);
              }
  
              if (StringUtils.isNotEmpty(name)) {
                  assertEquals(AtlasGraphUtilsV2.getEncodedProperty(v, ASSERT_NAME_PROPERTY, String.class), name, name);
              }
  
              count++;
          }
  
          assertEquals(count, expectedItems, String.format("%s:%s", typeName, name));
      }
  
      protected void assertMigrationStatus(int expectedTotalCount) {
          AtlasVertex v = getVertex("__MigrationStatus", "");
  
          assertTrue(AtlasGraphUtilsV2.getEncodedProperty(v, "currentIndex", Number.class).intValue() >= expectedTotalCount);
      }
  }
  