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;
  
  import com.google.common.collect.ImmutableList;
  import org.apache.atlas.exception.AtlasBaseException;
  import org.apache.atlas.glossary.GlossaryService;
  import org.apache.atlas.model.discovery.SearchParameters;
  import org.apache.atlas.model.glossary.AtlasGlossary;
  import org.apache.atlas.model.glossary.AtlasGlossaryTerm;
  import org.apache.atlas.model.glossary.relations.AtlasGlossaryHeader;
  import org.apache.atlas.model.instance.AtlasClassification;
  import org.apache.atlas.model.instance.AtlasEntity;
  import org.apache.atlas.model.instance.AtlasEntityHeader;
  import org.apache.atlas.model.instance.AtlasObjectId;
  import org.apache.atlas.model.instance.AtlasRelatedObjectId;
  import org.apache.atlas.model.instance.AtlasRelationship;
  import org.apache.atlas.model.instance.EntityMutationResponse;
  import org.apache.atlas.model.typedef.AtlasClassificationDef;
  import org.apache.atlas.model.typedef.AtlasEntityDef;
  import org.apache.atlas.model.typedef.AtlasEnumDef;
  import org.apache.atlas.model.typedef.AtlasStructDef;
  import org.apache.atlas.model.typedef.AtlasTypesDef;
  import org.apache.atlas.repository.AtlasTestBase;
  import org.apache.atlas.repository.store.graph.AtlasEntityStore;
  import org.apache.atlas.repository.store.graph.AtlasRelationshipStore;
  import org.apache.atlas.repository.store.graph.v2.AtlasEntityStream;
  import org.apache.atlas.store.AtlasTypeDefStore;
  import org.apache.atlas.type.AtlasTypeRegistry;
  import org.apache.atlas.utils.TestLoadModelUtils;
  import org.apache.commons.lang.StringUtils;
  
  import javax.inject.Inject;
  
  import java.io.IOException;
  import java.time.LocalDate;
  import java.time.ZoneId;
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Collections;
  import java.util.Date;
  import java.util.HashMap;
  import java.util.HashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Set;
  import java.util.stream.Collectors;
  import java.util.stream.Stream;
  
  import static org.apache.atlas.utils.TestLoadModelUtils.loadModelFromJson;
  import static org.testng.Assert.fail;
  
  public abstract class BasicTestSetup extends AtlasTestBase {
      //Classification type //
      public static final String DIMENSION_CLASSIFICATION   = "Dimension";
      public static final String FACT_CLASSIFICATION        = "Fact";
      public static final String PII_CLASSIFICATION         = "PII";
      public static final String METRIC_CLASSIFICATION      = "Metric";
      public static final String ETL_CLASSIFICATION         = "ETL";
      public static final String JDBC_CLASSIFICATION        = "JdbcAccess";
      public static final String LOGDATA_CLASSIFICATION     = "Log Data";
      public static final String DIMENSIONAL_CLASSIFICATION = "Dimensional";
  
      // Glossary type //
      public static final String SALES_GLOSSARY   = "salesGlossary";
      public static final String SALES_TERM       = "salesTerm";
      public static final String MODERNTRADE_TERM = "modernTrade";
      public static final String ECOMMERCE_TERM   = "ecommerce";
      public static final String USER_TYPE           = "user";
      public static final String POST_TYPE           = "post";
      public static final String HIGHLIGHT_TYPE      = "highlight";
      public static final String USER_POST_TYPE      = "user_post";
      public static final String HIGHLIGHT_POST_TYPE = "highlight_post";
  
      // Entity type //
      protected static final String DATABASE_TYPE     = "hive_db";
      protected static final String HIVE_TABLE_TYPE   = "hive_table";
      protected static final String DATASET_SUBTYPE   = "Asset";
      protected static final String HDFS_PATH         = "hdfs_path";
  
      private static final String COLUMN_TYPE       = "hive_column";
      private static final String HIVE_PROCESS_TYPE = "hive_process";
      private static final String STORAGE_DESC_TYPE = "hive_storagedesc";
      private static final String VIEW_TYPE         = "hive_process";
  
      @Inject
      protected AtlasTypeRegistry typeRegistry;
  
      @Inject
      protected AtlasTypeDefStore typeDefStore;
  
      @Inject
      protected AtlasEntityStore entityStore;
  
      @Inject
      protected AtlasRelationshipStore relationshipStore;
  
      @Inject
      protected GlossaryService glossaryService;
  
      private boolean                baseLoaded;
      private EntityMutationResponse hiveEntities;
      private AtlasEntity            timeDim;
      private AtlasEntity            productDim;
      private AtlasEntity            loggingFactMonthly;
  
      public AtlasEntity.AtlasEntitiesWithExtInfo hiveTestEntities() {
          List<AtlasEntity> entities = new ArrayList<>();
  
          createClassificationTypes();
  
          AtlasEntity salesDB = database("Sales", "/apps/warehouse/Sales Database", "John ETL", "hdfs://host:8000/apps/warehouse/sales");
          entities.add(salesDB);
  
          AtlasEntity sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));
  
          entities.add(sd);
  
          List<AtlasEntity> salesFactColumns = ImmutableList.of(column("time_id", "int", "time id"),
                  column("product_id", "int", "product id"),
                  column("customer_id", "int", "customer id", PII_CLASSIFICATION),
                  column("sales", "double", "product id", METRIC_CLASSIFICATION),
                  column("test", "int", "test 1"),
                  column("test_limit", "int", "test limit 1"));
  
          entities.addAll(salesFactColumns);
  
          AtlasEntity salesFact = table("sales_fact", "sales fact table", salesDB, sd, "Joe", "Managed", salesFactColumns, FACT_CLASSIFICATION);
  
          salesFact.setAttribute("createTime", new Date(2018, 01, 01));
  
          entities.add(salesFact);
  
          List<AtlasEntity> logFactColumns = ImmutableList.of(column("time_id", "int", "time id"),
                  column("app_id", "int", "app id"),
                  column("machine_id", "int", "machine id"),
                  column("log", "string", "log data", LOGDATA_CLASSIFICATION),
                  column("test", "int", "test 2"),
                  column("test_limit", "int", "test limit 2"));
  
          entities.addAll(logFactColumns);
  
          List<AtlasEntity> timeDimColumns = ImmutableList.of(column("time_id", "int", "time id"),
                  column("dayOfYear", "int", "day Of Year"),
                  column("weekDay", "int", "week Day"));
  
          entities.addAll(timeDimColumns);
  
          sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));
  
          entities.add(sd);
  
          timeDim = table("time_dim", "time dimension table", salesDB, sd, "John Doe", "External", timeDimColumns, DIMENSION_CLASSIFICATION);
  
          entities.add(timeDim);
  
          AtlasEntity reportingDB = database("Reporting", "/apps/warehouse/reporting database", "Jane BI", "hdfs://host:8000/apps/warehouse/reporting");
  
          entities.add(reportingDB);
  
          sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));
  
          entities.add(sd);
  
          AtlasEntity salesFactDaily = table("sales_fact_daily_mv", "sales fact daily materialized view", reportingDB, sd, "Joe BI", "Managed", salesFactColumns, METRIC_CLASSIFICATION);
  
          salesFactDaily.setAttribute("createTime", Date.from(LocalDate.of(2016, 8, 19).atStartOfDay(ZoneId.systemDefault()).toInstant()));
          entities.add(salesFactDaily);
  
          sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));
          entities.add(sd);
  
          AtlasEntity circularLineageTable1 = table("table1", "", reportingDB, sd, "Vimal", "Managed", salesFactColumns, METRIC_CLASSIFICATION);
          entities.add(circularLineageTable1);
  
          sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));
  
          entities.add(sd);
  
          AtlasEntity circularLineageTable2 = table("table2", "", reportingDB, sd, "Vimal 2", "Managed", salesFactColumns, METRIC_CLASSIFICATION);
  
          entities.add(circularLineageTable2);
  
          AtlasEntity circularLineage1Process = loadProcess("circularLineage1", "hive query for daily summary", "John ETL", ImmutableList.of(circularLineageTable1),
                  ImmutableList.of(circularLineageTable2), "create table as select ", "plan", "id", "graph", ETL_CLASSIFICATION);
  
          entities.add(circularLineage1Process);
  
          AtlasEntity circularLineage2Process = loadProcess("circularLineage2", "hive query for daily summary", "John ETL", ImmutableList.of(circularLineageTable2),
                  ImmutableList.of(circularLineageTable1), "create table as select ", "plan", "id", "graph", ETL_CLASSIFICATION);
  
          entities.add(circularLineage2Process);
  
          AtlasEntity loadSalesDaily = loadProcess("loadSalesDaily", "hive query for daily summary", "John ETL", ImmutableList.of(salesFact, timeDim),
                  ImmutableList.of(salesFactDaily), "create table as select ", "plan", "id", "graph", ETL_CLASSIFICATION);
  
          entities.add(loadSalesDaily);
  
          AtlasEntity logDB = database("Logging", "/apps/warehouse/logging", "Tim ETL", "hdfs://host:8000/apps/warehouse/logging");
  
          entities.add(logDB);
  
          sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));
  
          entities.add(sd);
  
          AtlasEntity loggingFactDaily = table("log_fact_daily_mv", "log fact daily materialized view", logDB, sd, "Tim ETL", "Managed", logFactColumns, LOGDATA_CLASSIFICATION);
  
          entities.add(loggingFactDaily);
  
          List<AtlasEntity> productDimColumns = ImmutableList.of(column("product_id", "int", "product id"),
                  column("product_name", "string", "product name"),
                  column("brand_name", "int", "brand name"));
  
          entities.addAll(productDimColumns);
  
          sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));
  
          entities.add(sd);
  
          productDim = table("product_dim", "product dimension table", salesDB, sd, "John Doe 2", "Managed", productDimColumns, DIMENSION_CLASSIFICATION);
  
          entities.add(productDim);
  
          AtlasEntity productDimView = view("product_dim_view", reportingDB, ImmutableList.of(productDim), DIMENSION_CLASSIFICATION, JDBC_CLASSIFICATION);
  
          entities.add(productDimView);
  
          List<AtlasEntity> customerDimColumns = ImmutableList.of(column("customer_id", "int", "customer id", PII_CLASSIFICATION),
                  column("name", "string", "customer name", PII_CLASSIFICATION),
                  column("address", "string", "customer address", PII_CLASSIFICATION));
  
          entities.addAll(customerDimColumns);
  
          sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));
  
          entities.add(sd);
  
          AtlasEntity customerDim = table("customer_dim", "customer dimension table", salesDB, sd, "fetl", "External", customerDimColumns, DIMENSION_CLASSIFICATION);
  
          entities.add(customerDim);
  
          AtlasEntity customerDimView = view("customer_dim_view", reportingDB, ImmutableList.of(customerDim), DIMENSION_CLASSIFICATION, JDBC_CLASSIFICATION);
  
          entities.add(customerDimView);
  
          sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));
  
          entities.add(sd);
  
          AtlasEntity salesFactMonthly = table("sales_fact_monthly_mv", "sales fact monthly materialized view", reportingDB, sd, "Jane BI", "Managed", salesFactColumns, METRIC_CLASSIFICATION);
  
          entities.add(salesFactMonthly);
  
          AtlasEntity loadSalesMonthly = loadProcess("loadSalesMonthly", "hive query for monthly summary", "John ETL", ImmutableList.of(salesFactDaily),
                  ImmutableList.of(salesFactMonthly), "create table as select ", "plan", "id", "graph", ETL_CLASSIFICATION);
  
          entities.add(loadSalesMonthly);
  
          sd = storageDescriptor("hdfs://host:8000/apps/warehouse/sales", "TextInputFormat", "TextOutputFormat", true, ImmutableList.of(column("time_id", "int", "time id")));
  
          entities.add(sd);
  
          loggingFactMonthly = table("logging_fact_monthly_mv", "logging fact monthly materialized view", logDB, sd, "Tim ETL 2", "Managed", logFactColumns, LOGDATA_CLASSIFICATION);
  
          entities.add(loggingFactMonthly);
  
          AtlasEntity loadLogsMonthly = loadProcess("loadLogsMonthly", "hive query for monthly summary", "Tim ETL", ImmutableList.of(loggingFactDaily),
                  ImmutableList.of(loggingFactMonthly), "create table as select ", "plan", "id", "graph", ETL_CLASSIFICATION);
  
          entities.add(loadLogsMonthly);
  
          AtlasEntity datasetSubType = datasetSubType("dataSetSubTypeInst1", "testOwner");
  
          entities.add(datasetSubType);
  
          return new AtlasEntity.AtlasEntitiesWithExtInfo(entities);
      }
  
      public EntityMutationResponse createDummyEntity(String name, String type, String... traitNames) throws AtlasBaseException {
          AtlasEntity entity = new AtlasEntity(type);
  
          entity.setAttribute("name", name);
          entity.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
          entity.setClassifications(Stream.of(traitNames).map(AtlasClassification::new).collect(Collectors.toList()));
  
          return entityStore.createOrUpdate(new AtlasEntityStream(new AtlasEntity.AtlasEntitiesWithExtInfo(entity)), false);
      }
  
      public SearchParameters.FilterCriteria getSingleFilterCondition(String attName, SearchParameters.Operator op, String attrValue) {
          SearchParameters.FilterCriteria filterCriteria = new SearchParameters.FilterCriteria();
          filterCriteria.setCondition(SearchParameters.FilterCriteria.Condition.AND);
  
          List<SearchParameters.FilterCriteria> criteria = new ArrayList<>();
          SearchParameters.FilterCriteria       f1       = new SearchParameters.FilterCriteria();
          f1.setAttributeName(attName);
          f1.setOperator(op);
          f1.setAttributeValue(attrValue);
          criteria.add(f1);
  
          filterCriteria.setCriterion(criteria);
          return filterCriteria;
      }
  
      public void assignGlossary() {
          try {
              AtlasGlossary glossary = new AtlasGlossary();
              glossary.setName(SALES_GLOSSARY);
              glossary = glossaryService.createGlossary(glossary);
  
              AtlasGlossaryTerm term = new AtlasGlossaryTerm();
              term.setAnchor(new AtlasGlossaryHeader(glossary.getGuid()));
              term.setName(SALES_TERM);
              term = glossaryService.createTerm(term);
  
              AtlasGlossaryTerm modernTrade = new AtlasGlossaryTerm();
              modernTrade.setAnchor(new AtlasGlossaryHeader(glossary.getGuid()));
              modernTrade.setName(MODERNTRADE_TERM);
              modernTrade = glossaryService.createTerm(modernTrade);
  
              AtlasGlossaryTerm ecomm = new AtlasGlossaryTerm();
              ecomm.setAnchor(new AtlasGlossaryHeader(glossary.getGuid()));
              ecomm.setName(ECOMMERCE_TERM);
              ecomm = glossaryService.createTerm(ecomm);
  
              List<AtlasRelatedObjectId> guids             = new ArrayList<>();
              List<AtlasRelatedObjectId> mordernTradeGuids = new ArrayList<>();
              List<AtlasRelatedObjectId> ecomGuid          = new ArrayList<>();
  
              for (AtlasEntityHeader p : hiveEntities.getCreatedEntities()) {
                  if (StringUtils.equals(p.getTypeName(), HIVE_TABLE_TYPE)) {
                      AtlasRelatedObjectId obj = new AtlasRelatedObjectId();
                      obj.setGuid(p.getGuid());
                      obj.setTypeName(p.getTypeName());
                      guids.add(obj);
  
                      if (p.getAttribute("name").equals(timeDim.getAttribute("name"))) {
                          timeDim.setGuid(p.getGuid());
                          AtlasRelatedObjectId obj1 = new AtlasRelatedObjectId();
                          obj1.setGuid(p.getGuid());
                          obj1.setTypeName(p.getTypeName());
                          mordernTradeGuids.add(obj1);
                      } else if (p.getAttribute("name").equals(productDim.getAttribute("name"))) {
                          productDim.setGuid(p.getGuid());
                          AtlasRelatedObjectId obj2 = new AtlasRelatedObjectId();
                          obj2.setGuid(p.getGuid());
                          obj2.setTypeName(p.getTypeName());
                          ecomGuid.add(obj2);
                      } else if (p.getAttribute("name").equals(loggingFactMonthly.getAttribute("name"))) {
                          loggingFactMonthly.setGuid(p.getGuid());
                          AtlasRelatedObjectId obj3 = new AtlasRelatedObjectId();
                          obj3.setGuid(p.getGuid());
                          obj3.setTypeName(p.getTypeName());
                          mordernTradeGuids.add(obj3);
                      }
                  }
              }
  
              glossaryService.assignTermToEntities(term.getGuid(), guids);
              glossaryService.assignTermToEntities(modernTrade.getGuid(), mordernTradeGuids);
              glossaryService.assignTermToEntities(ecomm.getGuid(), ecomGuid);
          } catch (AtlasBaseException e) {
              fail("Failed to assign glossary term");
          }
      }
  
      public void addLabels() {
          for (AtlasEntityHeader entity : hiveEntities.getCreatedEntities()) {
              if (entity.getTypeName().equals(HIVE_TABLE_TYPE)) {
                  Set<String> labels = new HashSet<>();
  
                  labels.add("你好");
  
                  try {
                      entityStore.setLabels(entity.getGuid(), labels);
                  } catch (AtlasBaseException e) {
                      fail("Failed to add Labels for Chinese characters to entity");
                  }
                  break;
              }
          }
      }
  
      protected void setupTestData() {
          loadBaseModels();
          loadHiveDataset();
          loadFsDataset();
          loadEmployeeDataset();
          assignGlossary();
      }
  
      protected void setupRelationshipTestData() {
          if (!baseLoaded) {
              loadBaseModels();
          }
  
          try {
              TestLoadModelUtils.loadModelFromResourcesJson("relationship_search_def.json", typeDefStore, typeRegistry);
          } catch (IOException e) {
              fail("Failed to load relationship_search_def.json");
          } catch (AtlasBaseException e) {
              fail("Failed to load relationship_search_def.json to atlas model");
          }
  
          try {
              createTestDataForRelationship();
          } catch (AtlasBaseException e) {
              fail("Failed to load data for relationship_search_def.json");
          }
      }
  
      protected void loadHiveDataset() {
          if (!baseLoaded) {
              loadBaseModels();
          }
  
          try {
              loadModelFromJson("1000-Hadoop/1030-hive_model.json", typeDefStore, typeRegistry);
          } catch (IOException | AtlasBaseException e) {
              fail("Hive model setup is required for test to run!");
          }
  
          AtlasEntity.AtlasEntitiesWithExtInfo hiveTestEntities = hiveTestEntities();
  
          try {
              hiveEntities = entityStore.createOrUpdate(new AtlasEntityStream(hiveTestEntities), false);
          } catch (AtlasBaseException e) {
              fail("Hive entities need to be created for test to run!");
          }
      }
  
      protected void loadFsDataset() {
          if (!baseLoaded) {
              loadBaseModels();
          }
  
          try {
              loadModelFromJson("1000-Hadoop/1020-fs_model.json", typeDefStore, typeRegistry);
          } catch (IOException | AtlasBaseException e) {
              fail("Fs model setup is required for test to run!");
          }
      }
  
      protected void loadEmployeeDataset() {
          if (!baseLoaded) {
              loadBaseModels();
          }
  
          // Define employee dataset types
          AtlasTypesDef employeeTypes = TestUtilsV2.defineDeptEmployeeTypes();
  
          try {
              typeDefStore.createTypesDef(employeeTypes);
          } catch (AtlasBaseException e) {
              fail("Employee Type setup is required");
          }
  
          // Define entities for department
          AtlasEntity.AtlasEntitiesWithExtInfo deptEg2 = TestUtilsV2.createDeptEg2();
  
          try {
              entityStore.createOrUpdate(new AtlasEntityStream(deptEg2), false);
          } catch (AtlasBaseException e) {
              fail("Employee entity setup should've passed");
          }
      }
  
      protected void createClassificationTypes() {
          List<AtlasClassificationDef> cds = Arrays.asList(new AtlasClassificationDef(DIMENSION_CLASSIFICATION, "Dimension Classification", "1.0",
                          Arrays.asList(new AtlasStructDef.AtlasAttributeDef("timeAttr", "string"), new AtlasStructDef.AtlasAttributeDef("productAttr", "string"))),
                  new AtlasClassificationDef(FACT_CLASSIFICATION, "Fact Classification", "1.0"),
                  new AtlasClassificationDef(PII_CLASSIFICATION, "PII Classification", "1.0"),
                  new AtlasClassificationDef(METRIC_CLASSIFICATION, "Metric Classification", "1.0"),
                  new AtlasClassificationDef(ETL_CLASSIFICATION, "ETL Classification", "1.0"),
                  new AtlasClassificationDef(JDBC_CLASSIFICATION, "JdbcAccess Classification", "1.0",
                          Arrays.asList(new AtlasStructDef.AtlasAttributeDef("attr1", "string"))),
                  new AtlasClassificationDef(LOGDATA_CLASSIFICATION, "LogData Classification", "1.0"),
                  new AtlasClassificationDef(DIMENSIONAL_CLASSIFICATION, "Dimensional Classification", "1.0",
                          Arrays.asList(new AtlasStructDef.AtlasAttributeDef("attr1", "string"))));
  
          AtlasTypesDef tds = new AtlasTypesDef(Collections.<AtlasEnumDef>emptyList(),
                  Collections.<AtlasStructDef>emptyList(),
                  cds,
                  Collections.<AtlasEntityDef>emptyList());
          createUpdateClassificationDef(tds);
      }
  
      protected AtlasEntity storageDescriptor(String location, String inputFormat, String outputFormat, boolean compressed, List<AtlasEntity> columns) {
          AtlasEntity storageDescriptor = new AtlasEntity(STORAGE_DESC_TYPE);
          storageDescriptor.setAttribute("location", location);
          storageDescriptor.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, "qualified:" + location);
          storageDescriptor.setAttribute("inputFormat", inputFormat);
          storageDescriptor.setAttribute("outputFormat", outputFormat);
          storageDescriptor.setAttribute("compressed", compressed);
          storageDescriptor.setAttribute("cols", getAtlasObjectIds(columns));
  
          return storageDescriptor;
      }
  
      protected AtlasEntity column(String name, String dataType, String comment, String... traitNames) {
          AtlasEntity column = new AtlasEntity(COLUMN_TYPE);
          column.setAttribute("name", name);
          column.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, "qualified:" + name);
          column.setAttribute("type", dataType);
          column.setAttribute("comment", comment);
          column.setClassifications(Stream.of(traitNames).map(AtlasClassification::new).collect(Collectors.toList()));
  
          return column;
      }
  
      protected AtlasEntity table(String name, String description, AtlasEntity db, AtlasEntity sd, String owner, String tableType,
              List<AtlasEntity> columns, String... traitNames) {
          String dbName      = db.getAttribute(AtlasClient.NAME).toString();
          String clusterName = db.getAttribute("clusterName").toString();
  
          AtlasEntity table = new AtlasEntity(HIVE_TABLE_TYPE);
          table.setAttribute("name", name);
          table.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, dbName + "." + name);
          table.setAttribute("description", description);
          table.setAttribute("owner", owner);
          table.setAttribute("tableType", tableType);
          table.setAttribute("temporary", false);
          table.setAttribute("createTime", new Date(System.currentTimeMillis()));
          table.setAttribute("lastAccessTime", System.currentTimeMillis());
          table.setAttribute("retention", 5d);
  
          table.setAttribute("db", getAtlasObjectId(db));
          table.setAttribute("sd", getAtlasObjectId(sd));
  
          table.setAttribute("columns", getAtlasObjectIds(columns));
          if (name.equals("time_dim")) {
              AtlasClassification classification = new AtlasClassification(traitNames[0], Collections.singletonMap("timeAttr", "timeValue"));
  
              table.setClassifications(Collections.singletonList(classification));
          } else if (name.equals("product_dim")) {
              AtlasClassification classification = new AtlasClassification(traitNames[0], Collections.singletonMap("productAttr", "productValue"));
  
              table.setClassifications(Collections.singletonList(classification));
          } else {
              table.setClassifications(Stream.of(traitNames).map(AtlasClassification::new).collect(Collectors.toList()));
          }
  
          sd.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, dbName + "." + name + "@" + clusterName + "_storage");
  
          AtlasObjectId tableId = getAtlasObjectId(table);
  
          sd.setAttribute("table", tableId);
  
          for (AtlasEntity column : columns) {
              column.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, dbName + "." + name + "." + column.getAttribute(AtlasClient.NAME).toString() + "@" + clusterName);
              column.setAttribute("table", tableId);
          }
  
          return table;
      }
  
      protected AtlasEntity loadProcess(String name, String description, String user,
              List<AtlasEntity> inputTables, List<AtlasEntity> outputTables,
              String queryText, String queryPlan, String queryId, String queryGraph, String... traitNames) {
          AtlasEntity process = new AtlasEntity(HIVE_PROCESS_TYPE);
          process.setAttribute("name", name);
          process.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
          process.setAttribute("description", description);
          process.setAttribute("userName", user);
          process.setAttribute("startTime", System.currentTimeMillis());
          process.setAttribute("endTime", System.currentTimeMillis() + 10000);
  
          process.setAttribute("operationType", "load");
          process.setAttribute("inputs", getAtlasObjectIds(inputTables));
          process.setAttribute("outputs", getAtlasObjectIds(outputTables));
  
          process.setAttribute("queryText", queryText);
          process.setAttribute("queryPlan", queryPlan);
          process.setAttribute("queryId", queryId);
          process.setAttribute("queryGraph", queryGraph);
  
          process.setClassifications(Stream.of(traitNames).map(AtlasClassification::new).collect(Collectors.toList()));
  
          return process;
      }
  
      AtlasEntity database(String name, String description, String owner, String locationUri, String... traitNames) {
          AtlasEntity database = new AtlasEntity(DATABASE_TYPE);
          database.setAttribute("name", name);
          database.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, "qualified:" + name);
          database.setAttribute("description", description);
          database.setAttribute("userDescription", description);
          database.setAttribute("owner", owner);
          database.setAttribute("location", locationUri);
          database.setAttribute("createTime", System.currentTimeMillis());
          database.setAttribute("clusterName", "cl1");
          database.setClassifications(Stream.of(traitNames).map(AtlasClassification::new).collect(Collectors.toList()));
  
          return database;
      }
  
      AtlasEntity view(String name, AtlasEntity dbId, List<AtlasEntity> inputTables, String... traitNames) {
          AtlasEntity view = new AtlasEntity(VIEW_TYPE);
          view.setAttribute("name", name);
          view.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
          view.setAttribute("userName", "testUser");
          view.setAttribute("startTime", System.currentTimeMillis());
          view.setAttribute("endTime", System.currentTimeMillis() + 10000);
  
          view.setAttribute("operationType", "view");
          view.setAttribute("query", "create table as select");
          view.setAttribute("queryText", "create table as select");
          view.setAttribute("queryPlan", "viewPlan");
          view.setAttribute("queryId", "view1");
          view.setAttribute("db", getAtlasObjectId(dbId));
          view.setAttribute("inputs", getAtlasObjectIds(inputTables));
          view.setClassifications(Stream.of(traitNames).map(AtlasClassification::new).collect(Collectors.toList()));
  
          return view;
      }
  
      AtlasEntity datasetSubType(final String name, String owner) {
          AtlasEntity datasetSubType = new AtlasEntity(DATASET_SUBTYPE);
          datasetSubType.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, name);
          datasetSubType.setAttribute(AtlasClient.NAME, name);
          datasetSubType.setAttribute("owner", owner);
  
          return datasetSubType;
      }
  
      private void loadBaseModels() {
          try {
              loadModelFromJson("0000-Area0/0010-base_model.json", typeDefStore, typeRegistry);
              loadModelFromJson("0000-Area0/0011-glossary_model.json", typeDefStore, typeRegistry);
              baseLoaded = true;
          } catch (IOException | AtlasBaseException e) {
              fail("Base model setup is required for test to run!");
          }
      }
  
      private void createUpdateClassificationDef(AtlasTypesDef td) {
          try {
              typeDefStore.createTypesDef(td);
          } catch (Exception e) {
              fail("Error creating classifications definitions.");
          }
      }
  
      private List<AtlasObjectId> getAtlasObjectIds(List<AtlasEntity> columns) {
          List<AtlasObjectId> objIds = new ArrayList<>();
          for (AtlasEntity e : columns) {
              AtlasObjectId oid = getAtlasObjectId(e);
              objIds.add(oid);
          }
          return objIds;
      }
  
      private AtlasObjectId getAtlasObjectId(AtlasEntity e) {
          return new AtlasObjectId(e.getGuid(), e.getTypeName());
      }
  
      private void createTestDataForRelationship() throws AtlasBaseException {
          List<AtlasEntity> entities = new ArrayList<>();
  
          //users
          AtlasEntity ajayUser = createUser("Ajay Mishra", new Date(1985, 10, 9), "Male", "Content Writer");
          entities.add(ajayUser);
          AtlasEntity maryUser = createUser("Mary Williams", new Date(1998, 2, 19), "Female", "Architecture");
          entities.add(maryUser);
          AtlasEntity divyaUser = createUser("Divya Deshmukh", new Date(1980, 12, 10), "Female", "Designer");
          entities.add(divyaUser);
  
          //posts
          AtlasEntity indAugPost = createPost("ind-post@Ajay", "Image", new Date(2021, 8, 15), "Independence Day Celebration");
          entities.add(indAugPost);
          AtlasEntity onamAugPost = createPost("onam-post@Ajay", "Image", new Date(2021, 8, 8), "Onam with Family");
          entities.add(onamAugPost);
          AtlasEntity diwaliPost = createPost("diwali-post@Ajay", "Reel", new Date(2020, 11, 14), "Diwali get together");
          entities.add(diwaliPost);
          AtlasEntity tripPost = createPost("trip-post@Ajay", "Video", new Date(2021, 3, 10), "Kashmir a serene");
          entities.add(tripPost);
          AtlasEntity officeParty = createPost("officeparty-post@Ajay", "Video", new Date(2021, 4, 4), "office party");
          entities.add(officeParty);
  
          AtlasEntity indPost = createPost("ind-post@Mary", "Video", new Date(2021, 8, 15), "Independence Day Celebration");
          entities.add(indPost);
          AtlasEntity christPost = createPost("christmas-post@Mary", "Image", new Date(2020, 12, 25), "Christmas Day");
          entities.add(christPost);
          AtlasEntity sraPost = createPost("elevationSRA@Mary", "Image", new Date(2020, 12, 25), "Elevation");
          entities.add(sraPost);
  
          AtlasEntity ganeshPost = createPost("ganeshchaturthi-post@Divya", "Reel", new Date(2020, 12, 25), "Ganesh chaturthi");
          entities.add(ganeshPost);
  
          //highlights
          AtlasEntity year2021 = createHighlight("year2021@Ajay", "Journey of 2021", new Date(2021, 9, 25));
          entities.add(year2021);
          AtlasEntity festival = createHighlight("festives@Ajay", "Festive Celebration", new Date(2021, 9, 25));
          entities.add(festival);
  
          AtlasEntity projectSRA = createHighlight("projectSRA@Mary", "Project SRA", new Date(2021, 9, 25));
          entities.add(projectSRA);
  
          // create entity types
          EntityMutationResponse  response = entityStore.createOrUpdate(new AtlasEntityStream(entities), false);
          List<AtlasEntityHeader> headers  = response.getCreatedEntities();
          setGuidInHeaders(entities, headers);
  
          // create relationship types
          createUserPostRelationship(ajayUser, indAugPost, "create");
          createUserPostRelationship(ajayUser, onamAugPost, "create");
          createUserPostRelationship(ajayUser, diwaliPost, "create");
          createUserPostRelationship(ajayUser, tripPost, "create");
          createUserPostRelationship(ajayUser, officeParty, "create");
          createUserPostRelationship(ajayUser, christPost, "like");
          createUserPostRelationship(ajayUser, ganeshPost, "wow");
  
          createUserPostRelationship(maryUser, indPost, "create");
          createUserPostRelationship(maryUser, christPost, "create");
          createUserPostRelationship(maryUser, sraPost, "create");
          createUserPostRelationship(maryUser, onamAugPost, "wow");
  
          createUserPostRelationship(divyaUser, ganeshPost, "create");
          createUserPostRelationship(divyaUser, indAugPost, "like");
          createUserPostRelationship(divyaUser, indPost, "like");
          createUserPostRelationship(divyaUser, onamAugPost, "love");
          createUserPostRelationship(divyaUser, diwaliPost, "love");
          createUserPostRelationship(divyaUser, christPost, "like");
  
          createHighlightPostRelationship(year2021, indAugPost);
          createHighlightPostRelationship(year2021, onamAugPost);
          createHighlightPostRelationship(year2021, tripPost);
          createHighlightPostRelationship(year2021, officeParty);
  
          createHighlightPostRelationship(festival, indAugPost);
          createHighlightPostRelationship(festival, onamAugPost);
          createHighlightPostRelationship(festival, diwaliPost);
  
          createHighlightPostRelationship(projectSRA, sraPost);
      }
  
      private AtlasEntity createUser(String id, Date birthDate, String gender, String bio) {
          AtlasEntity user = new AtlasEntity(USER_TYPE);
          user.setAttribute("id", id);
          user.setAttribute("birthDate", birthDate);
          user.setAttribute("gender", gender);
          user.setAttribute("bio", bio);
  
          return user;
      }
  
      private AtlasEntity createPost(String id, String type, Date postDate, String caption) {
          AtlasEntity post = new AtlasEntity(POST_TYPE);
          post.setAttribute("type", type);
          post.setAttribute("post_date", postDate);
          post.setAttribute("caption", caption);
          post.setAttribute("id", id);
  
          return post;
      }
  
      private AtlasEntity createHighlight(String id, String title, Date createdOn) {
          AtlasEntity highlight = new AtlasEntity(HIGHLIGHT_TYPE);
          highlight.setAttribute("created_on", createdOn);
          highlight.setAttribute("title", title);
          highlight.setAttribute("id", id);
  
          return highlight;
      }
  
      private void createUserPostRelationship(AtlasEntity user, AtlasEntity post, String reaction) {
          Map<String, Object> attributes = new HashMap<>();
          attributes.put("user_name", user.getAttribute("id"));
          attributes.put("post_name", post.getAttribute("id"));
          attributes.put("reaction", reaction);
          try {
              AtlasRelationship relationship = new AtlasRelationship(USER_POST_TYPE, getAtlasObjectId(user), getAtlasObjectId(post), attributes);
              relationshipStore.create(relationship);
          } catch (AtlasBaseException e) {
              fail("Failed to create Relationship between , " + user.getTypeName() + ", " + post.getTypeName());
          }
      }
  
      private void createHighlightPostRelationship(AtlasEntity highlight, AtlasEntity post) {
          Map<String, Object> attributes = new HashMap<>();
          attributes.put("highlight_name", highlight.getAttribute("id"));
          attributes.put("post_name", post.getAttribute("id"));
          try {
              AtlasRelationship relationship = new AtlasRelationship(HIGHLIGHT_POST_TYPE, getAtlasObjectId(highlight), getAtlasObjectId(post), attributes);
              relationshipStore.create(relationship);
          } catch (AtlasBaseException e) {
              fail("Failed to create Relationship between , " + highlight.getTypeName() + ", " + post.getTypeName());
          }
      }
  
      private List<AtlasEntity> setGuidInHeaders(List<AtlasEntity> entities, List<AtlasEntityHeader> createdEntities) {
          List<String> assignedGuids = new ArrayList<>();
  
          for (AtlasEntityHeader header : createdEntities) {
              String id = header.getAttribute("id").toString();
  
              for (AtlasEntity entity : entities) {
                  if (id.equalsIgnoreCase(entity.getAttribute("id").toString())) {
                      String guid = header.getGuid();
                      entity.setGuid(guid);
                      assignedGuids.add(guid);
  
                      break;
                  }
              }
          }
          return entities;
      }
  }
  