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.web.integration;
  
  import com.fasterxml.jackson.core.type.TypeReference;
  import com.fasterxml.jackson.databind.ObjectMapper;
  import com.fasterxml.jackson.databind.node.ArrayNode;
  import com.google.common.collect.Lists;
  import com.sun.jersey.api.client.ClientResponse;
  import org.apache.atlas.ApplicationProperties;
  import org.apache.atlas.AtlasClient;
  import org.apache.atlas.AtlasServiceException;
  import org.apache.atlas.EntityAuditEvent;
  import org.apache.atlas.bulkimport.BulkImportResponse;
  import org.apache.atlas.model.TimeBoundary;
  import org.apache.atlas.model.audit.AtlasAuditEntry;
  import org.apache.atlas.model.audit.AuditSearchParameters;
  import org.apache.atlas.model.audit.EntityAuditEventV2;
  import org.apache.atlas.model.instance.AtlasClassification;
  import org.apache.atlas.model.instance.AtlasClassification.AtlasClassifications;
  import org.apache.atlas.model.instance.AtlasEntity;
  import org.apache.atlas.model.instance.AtlasEntity.AtlasEntityWithExtInfo;
  import org.apache.atlas.model.instance.AtlasEntityHeader;
  import org.apache.atlas.model.instance.AtlasObjectId;
  import org.apache.atlas.model.instance.ClassificationAssociateRequest;
  import org.apache.atlas.model.instance.EntityMutationResponse;
  import org.apache.atlas.model.instance.EntityMutations;
  import org.apache.atlas.model.typedef.AtlasClassificationDef;
  import org.apache.atlas.model.typedef.AtlasEntityDef;
  import org.apache.atlas.model.typedef.AtlasTypesDef;
  import org.apache.atlas.type.AtlasTypeUtil;
  import org.apache.atlas.utils.TestResourceFileUtils;
  import org.apache.atlas.v1.typesystem.types.utils.TypesUtil;
  import org.apache.commons.io.FileUtils;
  import org.apache.commons.lang.RandomStringUtils;
  import org.apache.hadoop.util.StringUtils;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  import org.testng.annotations.BeforeClass;
  import org.testng.annotations.Test;
  
  import java.io.File;
  import java.io.IOException;
  import java.nio.charset.Charset;
  import java.nio.charset.StandardCharsets;
  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.testng.Assert.assertEquals;
  import static org.testng.Assert.assertFalse;
  import static org.testng.Assert.assertNotEquals;
  import static org.testng.Assert.assertNotNull;
  import static org.testng.Assert.assertNull;
  import static org.testng.Assert.assertTrue;
  import static org.testng.Assert.fail;
  
  /**
   * Integration tests for Entity Jersey Resource.
   */
  public class EntityV2JerseyResourceIT extends BaseResourceIT {
      private static final Logger LOG = LoggerFactory.getLogger(EntityV2JerseyResourceIT.class);
  
      private static final String ENTITY_NOTIFICATION_VERSION_PROPERTY = "atlas.notification.entity.version";
  
      private final String databaseName = "db" + randomString();
      private final String tableName    = "table" + randomString();
      private       String traitName;
      private       String createdDBName;
      private       String createdTableQualifiedName;
  
      private AtlasEntity            dbEntity;
      private AtlasEntity            tableEntity;
      private AtlasClassificationDef piiTrait;
  
      @BeforeClass
      public void setUp() throws Exception {
          super.setUp();
  
          createTypeDefinitionsV2();
      }
  
      @Test
      public void testSubmitEntity() throws Exception {
          TypesUtil.Pair dbAndTable = createDBAndTable();
          assertNotNull(dbAndTable);
          assertNotNull(dbAndTable.left);
          assertNotNull(dbAndTable.right);
          // Writing created table data to a file for import test.
          createImportFile();
      }
  
      @Test
      public void testCreateNestedEntities() throws Exception {
          AtlasEntity.AtlasEntitiesWithExtInfo entities = new AtlasEntity.AtlasEntitiesWithExtInfo();
  
          AtlasEntity databaseInstance = new AtlasEntity(DATABASE_TYPE_V2, "name", "db1");
          databaseInstance.setAttribute("name", "db1");
          databaseInstance.setAttribute("description", "foo database");
          databaseInstance.setAttribute("owner", "user1");
          databaseInstance.setAttribute("locationUri", "/tmp");
          databaseInstance.setAttribute("createTime", 1000);
          entities.addEntity(databaseInstance);
  
          int nTables      = 5;
          int colsPerTable = 3;
  
          for (int i = 0; i < nTables; i++) {
              String tableName = "db1-table-" + i;
  
              AtlasEntity tableInstance = new AtlasEntity(HIVE_TABLE_TYPE_V2, "name", tableName);
              tableInstance.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, tableName);
              tableInstance.setAttribute("db", AtlasTypeUtil.getAtlasObjectId(databaseInstance));
              tableInstance.setAttribute("description", tableName + " table");
              entities.addEntity(tableInstance);
  
              List<AtlasObjectId> columns = new ArrayList<>();
              for (int j = 0; j < colsPerTable; j++) {
                  AtlasEntity columnInstance = new AtlasEntity(COLUMN_TYPE_V2);
                  columnInstance.setAttribute("name", tableName + "-col-" + j);
                  columnInstance.setAttribute("dataType", "String");
                  columnInstance.setAttribute("comment", "column " + j + " for table " + i);
  
                  columns.add(AtlasTypeUtil.getAtlasObjectId(columnInstance));
  
                  entities.addReferredEntity(columnInstance);
              }
              tableInstance.setAttribute("columns", columns);
          }
  
          //Create the tables.  The database and columns should be created automatically, since
          //the tables reference them.
  
          EntityMutationResponse response = atlasClientV2.createEntities(entities);
          assertNotNull(response);
  
          Map<String, String> guidsCreated = response.getGuidAssignments();
          assertEquals(guidsCreated.size(), nTables * colsPerTable + nTables + 1);
          assertNotNull(guidsCreated.get(databaseInstance.getGuid()));
  
          for (AtlasEntity r : entities.getEntities()) {
              assertNotNull(guidsCreated.get(r.getGuid()));
          }
  
          for (AtlasEntity r : entities.getReferredEntities().values()) {
              assertNotNull(guidsCreated.get(r.getGuid()));
          }
      }
  
      @Test
      public void testRequestUser() throws Exception {
          AtlasEntity            hiveDBInstanceV2 = createHiveDB(randomString());
          List<EntityAuditEvent> events           = atlasClientV1.getEntityAuditEvents(hiveDBInstanceV2.getGuid(), (short) 10);
          assertEquals(events.size(), 1);
          assertEquals(events.get(0).getUser(), "admin");
      }
  
      @Test
      public void testEntityDeduping() throws Exception {
          ArrayNode results = searchByDSL(String.format("%s where name='%s'", DATABASE_TYPE_V2, databaseName));
          assertEquals(results.size(), 1);
  
          final AtlasEntity hiveDBInstanceV2 = createHiveDB();
  
          results = searchByDSL(String.format("%s where name='%s'", DATABASE_TYPE_V2, databaseName));
          assertEquals(results.size(), 1);
  
          //Test the same across references
          final String tableName           = randomString();
          AtlasEntity  hiveTableInstanceV2 = createHiveTableInstanceV2(hiveDBInstanceV2, tableName);
          hiveTableInstanceV2.setAttribute(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, tableName);
  
          EntityMutationResponse entity = atlasClientV2.createEntity(new AtlasEntityWithExtInfo(hiveTableInstanceV2));
          assertNotNull(entity);
          assertNotNull(entity.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE));
          results = searchByDSL(String.format("%s where name='%s'", DATABASE_TYPE_V2, databaseName));
          assertEquals(results.size(), 1);
      }
  
      @Test
      public void testEntityDefinitionAcrossTypeUpdate() throws Exception {
          //create type
          AtlasEntityDef entityDef = AtlasTypeUtil.createClassTypeDef(randomString(), Collections.emptySet(), AtlasTypeUtil.createUniqueRequiredAttrDef("name", "string"));
          AtlasTypesDef typesDef = new AtlasTypesDef();
  
          typesDef.getEntityDefs().add(entityDef);
  
          AtlasTypesDef created = atlasClientV2.createAtlasTypeDefs(typesDef);
          assertNotNull(created);
          assertNotNull(created.getEntityDefs());
          assertEquals(created.getEntityDefs().size(), 1);
  
          //create entity for the type
          AtlasEntity instance = new AtlasEntity(entityDef.getName());
          instance.setAttribute("name", randomString());
          EntityMutationResponse mutationResponse = atlasClientV2.createEntity(new AtlasEntityWithExtInfo(instance));
          assertNotNull(mutationResponse);
          assertNotNull(mutationResponse.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE));
          assertEquals(mutationResponse.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).size(), 1);
          String guid = mutationResponse.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).get(0).getGuid();
  
          //update type - add attribute
          entityDef = AtlasTypeUtil.createClassTypeDef(entityDef.getName(), Collections.emptySet(),
                  AtlasTypeUtil.createUniqueRequiredAttrDef("name", "string"),
                  AtlasTypeUtil.createOptionalAttrDef("description", "string"));
  
          typesDef = new AtlasTypesDef();
          typesDef.getEntityDefs().add(entityDef);
  
          AtlasTypesDef updated = atlasClientV2.updateAtlasTypeDefs(typesDef);
          assertNotNull(updated);
          assertNotNull(updated.getEntityDefs());
          assertEquals(updated.getEntityDefs().size(), 1);
  
          //Get definition after type update - new attributes should be null
          AtlasEntity entityByGuid = getEntityByGuid(guid);
          assertNull(entityByGuid.getAttribute("description"));
          assertEquals(entityByGuid.getAttribute("name"), instance.getAttribute("name"));
      }
  
      @Test
      public void testEntityInvalidValue() {
          AtlasEntity databaseInstance = new AtlasEntity(DATABASE_TYPE_V2);
          String      dbName           = randomString();
          String      nullString       = null;
          String      emptyString      = "";
          databaseInstance.setAttribute("name", dbName);
          databaseInstance.setAttribute("description", nullString);
          AtlasEntityHeader created = createEntity(databaseInstance);
  
          // null valid value for required attr - description
          assertNull(created);
  
          databaseInstance.setAttribute("description", emptyString);
          created = createEntity(databaseInstance);
  
          // empty string valid value for required attr
          assertNotNull(created);
  
          databaseInstance.setGuid(created.getGuid());
          databaseInstance.setAttribute("owner", nullString);
          databaseInstance.setAttribute("locationUri", emptyString);
  
          created = updateEntity(databaseInstance);
  
          // null/empty string valid value for optional attr
          assertNotNull(created);
      }
  
      @Test
      public void testGetEntityByAttribute() throws Exception {
          AtlasEntity hiveDB        = createHiveDB();
          String      qualifiedName = (String) hiveDB.getAttribute(NAME);
          //get entity by attribute
  
          AtlasEntity byAttribute = atlasClientV2.getEntityByAttribute(DATABASE_TYPE_V2, toMap(NAME, qualifiedName)).getEntity();
          assertEquals(byAttribute.getTypeName(), DATABASE_TYPE_V2);
          assertEquals(byAttribute.getAttribute(NAME), qualifiedName);
      }
  
      @Test
      public void testGetEntitiesByAttribute() throws Exception {
          AtlasEntity hiveDB1 = createHiveDB();
          AtlasEntity hiveDB2 = createHiveDB();
  
          String                    qualifiedNameDB1 = (String) hiveDB1.getAttribute(NAME);
          String                    qualifiedNameDB2 = (String) hiveDB2.getAttribute(NAME);
          List<Map<String, String>> list             = new ArrayList<>();
          list.add(toMap(NAME, qualifiedNameDB1));
          list.add(toMap(NAME, qualifiedNameDB2));
  
          AtlasEntity.AtlasEntitiesWithExtInfo info       = atlasClientV2.getEntitiesByAttribute(DATABASE_TYPE_V2, list);
          List<AtlasEntity>                    entityList = info.getEntities();
          assertEquals(entityList.size(), 2);
          assertEquals(entityList.get(0).getTypeName(), DATABASE_TYPE_V2);
          assertEquals(entityList.get(1).getTypeName(), DATABASE_TYPE_V2);
      }
  
      @Test
      public void testSubmitEntityWithBadDateFormat() throws Exception {
          AtlasEntity       hiveDBEntity = createHiveDBInstanceV2("db" + randomString());
          AtlasEntityHeader hiveDBHeader = createEntity(hiveDBEntity);
          hiveDBEntity.setGuid(hiveDBHeader.getGuid());
  
          AtlasEntity tableInstance = createHiveTableInstanceV2(hiveDBEntity, "table" + randomString());
          //Dates with an invalid format are simply nulled out.  This does not produce
          //an error.  See AtlasBuiltInTypes.AtlasDateType.getNormalizedValue().
          tableInstance.setAttribute("lastAccessTime", 1107201407);
          AtlasEntityHeader tableEntityHeader = createEntity(tableInstance);
          assertNotNull(tableEntityHeader);
      }
  
      @Test(dependsOnMethods = "testSubmitEntity")
      public void testAddProperty() throws Exception {
          //add property
          String description = "bar table - new desc";
          addProperty(createHiveTable().getGuid(), "description", description);
  
          AtlasEntity entityByGuid = getEntityByGuid(createHiveTable().getGuid());
          assertNotNull(entityByGuid);
  
          entityByGuid.setAttribute("description", description);
  
          // TODO: This behavior should've been consistent across APIs
  //        //invalid property for the type
  //        try {
  //            addProperty(table.getGuid(), "invalid_property", "bar table");
  //            fail("Expected AtlasServiceException");
  //        } catch (AtlasServiceException e) {
  //            assertNotNull(e.getStatus());
  //            assertEquals(e.getStatus(), ClientResponse.Status.BAD_REQUEST);
  //        }
  
          //non-string property, update
          Object currentTime = new Date(System.currentTimeMillis());
          addProperty(createHiveTable().getGuid(), "createTime", currentTime);
  
          entityByGuid = getEntityByGuid(createHiveTable().getGuid());
          assertNotNull(entityByGuid);
      }
  
      @Test
      public void testAddNullPropertyValue() throws Exception {
          // FIXME: Behavior has changed between v1 and v2
          //add property
  //        try {
          addProperty(createHiveTable().getGuid(), "description", null);
  //            fail("Expected AtlasServiceException");
  //        } catch(AtlasServiceException e) {
  //            assertEquals(e.getStatus().getStatusCode(), Response.Status.BAD_REQUEST.getStatusCode());
  //        }
      }
  
      @Test(expectedExceptions = AtlasServiceException.class)
      public void testGetInvalidEntityDefinition() throws Exception {
          getEntityByGuid("blah");
      }
  
      @Test(dependsOnMethods = "testSubmitEntity", enabled = false)
      public void testGetEntityList() {
          // TODO: Can only be done when there's a search API exposed from entity REST
      }
  
      @Test(enabled = false)
      public void testGetEntityListForBadEntityType() {
          // FIXME: Complete test when search interface is in place
      }
  
      @Test(enabled = false)
      public void testGetEntityListForNoInstances() {
          // FIXME: Complete test when search interface is in place
          /*
          String typeName = "";
  
          ClientResponse clientResponse =
                  service.path(ENTITIES).queryParam("type", typeName).accept(Servlets.JSON_MEDIA_TYPE)
                          .type(Servlets.JSON_MEDIA_TYPE).method(HttpMethod.GET, ClientResponse.class);
          assertEquals(clientResponse.getStatus(), Response.Status.OK.getStatusCode());
  
          String responseAsString = clientResponse.getEntity(String.class);
          assertNotNull(responseAsString);
  
          JSONObject response = new JSONObject(responseAsString);
          assertNotNull(response.get(AtlasClient.REQUEST_ID));
  
          final JSONArray list = response.getJSONArray(AtlasClient.RESULTS);
          assertEquals(list.length(), 0);
           */
      }
  
      @Test(dependsOnMethods = "testSubmitEntity")
      public void testGetTraitNames() throws Exception {
          AtlasClassifications classifications = atlasClientV2.getClassifications(createHiveTable().getGuid());
          assertNotNull(classifications);
          assertFalse(classifications.getList().isEmpty());
          assertEquals(classifications.getList().size(), 9);
      }
  
      @Test(dependsOnMethods = "testSubmitEntity")
      public void testCommonAttributes() throws Exception {
          AtlasEntity entity = getEntityByGuid(createHiveTable().getGuid());
          assertNotNull(entity.getStatus());
          assertNotNull(entity.getVersion());
          assertNotNull(entity.getCreatedBy());
          assertNotNull(entity.getCreateTime());
          assertNotNull(entity.getUpdatedBy());
          assertNotNull(entity.getUpdateTime());
      }
  
      @Test(dependsOnMethods = "testGetTraitNames")
      public void testAddTrait() throws Exception {
          traitName = "PII_Trait" + randomString();
          AtlasClassificationDef piiTrait =
                  AtlasTypeUtil.createTraitTypeDef(traitName, Collections.emptySet());
          AtlasTypesDef typesDef = new AtlasTypesDef();
          typesDef.getClassificationDefs().add(piiTrait);
          createType(typesDef);
  
          atlasClientV2.addClassifications(createHiveTable().getGuid(), Collections.singletonList(new AtlasClassification(piiTrait.getName())));
  
          assertEntityAudit(createHiveTable().getGuid(), EntityAuditEvent.EntityAuditAction.TAG_ADD);
          AtlasClassifications classifications = atlasClientV2.getEntityClassifications(createHiveTable().getGuid(), piiTrait.getName());
          assertNotNull(classifications);
      }
  
      @Test(dependsOnMethods = "testAddTrait")
      public void testAddLabels() throws Exception {
          Set<String> set = new HashSet<>();
          set.add("lable");
          atlasClientV2.addLabels(createHiveTable().getGuid(), set);
          AtlasEntityWithExtInfo info = atlasClientV2.getEntityByGuid(createHiveTable().getGuid(), false, true);
          assertNotNull(info);
          assertNotNull(info.getEntity().getLabels());
          assertEquals(info.getEntity().getLabels().size(), 1);
      }
  
      @Test(dependsOnMethods = "testAddLabels")
      public void testSetLabels() throws Exception {
          Set<String> setNet = new HashSet<>();
          setNet.add("labelNext");
          atlasClientV2.setLabels(createHiveTable().getGuid(), setNet);
          AtlasEntityWithExtInfo infoForSet = atlasClientV2.getEntityByGuid(createHiveTable().getGuid(), false, true);
          assertNotNull(infoForSet);
          assertNotNull(infoForSet.getEntity().getLabels());
          assertEquals(infoForSet.getEntity().getLabels().size(), 1);
      }
  
      @Test(dependsOnMethods = "testSetLabels")
      public void testDeleteLabels() throws Exception {
          Set<String> set = new HashSet<>();
          set.add("testNext");
          atlasClientV2.removeLabels(createHiveTable().getGuid(), set);
          AtlasEntityWithExtInfo info = atlasClientV2.getEntityByGuid(createHiveTable().getGuid(), false, true);
          assertNotNull(info);
          assertNotNull(info.getEntity().getLabels());
          assertEquals(info.getEntity().getLabels().size(), 1);
      }
  
      @Test(dependsOnMethods = "testGetTraitNames")
      public void testAddTraitWithValidityPeriod() throws Exception {
          traitName = "PII_Trait" + randomString();
  
          AtlasClassificationDef piiTrait = AtlasTypeUtil.createTraitTypeDef(traitName, Collections.emptySet());
          AtlasTypesDef          typesDef = new AtlasTypesDef(Collections.emptyList(), Collections.emptyList(), Collections.singletonList(piiTrait), Collections.emptyList());
  
          createType(typesDef);
  
          String              tableGuid      = createHiveTable().getGuid();
          AtlasClassification classification = new AtlasClassification(piiTrait.getName());
          TimeBoundary        validityPeriod = new TimeBoundary("2018/03/01 00:00:00", "2018/04/01 00:00:00", "GMT");
  
          classification.setEntityGuid(tableGuid);
          classification.addValityPeriod(validityPeriod);
          classification.setPropagate(true);
          classification.setRemovePropagationsOnEntityDelete(true);
  
          atlasClientV2.addClassifications(tableGuid, Collections.singletonList(classification));
  
          assertEntityAudit(tableGuid, EntityAuditEvent.EntityAuditAction.TAG_ADD);
  
          AtlasClassifications classifications = atlasClientV2.getClassifications(tableGuid);
  
          assertNotNull(classifications);
          assertNotNull(classifications.getList());
          assertTrue(classifications.getList().size() > 1);
  
          boolean foundClassification = false;
          for (AtlasClassification entityClassification : classifications.getList()) {
              if (StringUtils.equalsIgnoreCase(entityClassification.getTypeName(), piiTrait.getName())) {
                  foundClassification = true;
  
                  assertEquals(entityClassification.getTypeName(), piiTrait.getName());
                  assertNotNull(entityClassification.getValidityPeriods());
                  assertEquals(entityClassification.getValidityPeriods().size(), 1);
                  assertEquals(entityClassification.getValidityPeriods().get(0), validityPeriod);
                  assertEquals(entityClassification, classification);
  
                  break;
              }
          }
  
          assertTrue(foundClassification, "classification '" + piiTrait.getName() + "' is missing for entity '" + tableGuid + "'");
      }
  
      @Test(dependsOnMethods = "testSubmitEntity")
      public void testGetTraitDefinitionForEntity() throws Exception {
          traitName = "PII_Trait" + randomString();
          AtlasClassificationDef piiTrait =
                  AtlasTypeUtil.createTraitTypeDef(traitName, Collections.emptySet());
          AtlasTypesDef typesDef = new AtlasTypesDef();
          typesDef.getClassificationDefs().add(piiTrait);
          createType(typesDef);
  
          AtlasClassificationDef classificationByName = atlasClientV2.getClassificationDefByName(traitName);
          assertNotNull(classificationByName);
  
          AtlasEntity hiveTable = createHiveTable();
          assertEquals(hiveTable.getClassifications().size(), 7);
  
          AtlasClassification piiClassification = new AtlasClassification(piiTrait.getName());
  
          atlasClientV2.addClassifications(hiveTable.getGuid(), Lists.newArrayList(piiClassification));
  
          AtlasClassifications classifications = atlasClientV2.getClassifications(hiveTable.getGuid());
          assertNotNull(classifications);
          assertFalse(classifications.getList().isEmpty());
          assertEquals(classifications.getList().size(), 9);
      }
  
      @Test(dependsOnMethods = "testGetTraitNames")
      public void testAddTraitWithAttribute() throws Exception {
          final String traitName = "PII_Trait" + randomString();
          AtlasClassificationDef piiTrait = AtlasTypeUtil
                  .createTraitTypeDef(traitName, Collections.emptySet(),
                          AtlasTypeUtil.createRequiredAttrDef("type", "string"));
          AtlasTypesDef typesDef = new AtlasTypesDef();
          typesDef.getClassificationDefs().add(piiTrait);
          createType(typesDef);
  
          AtlasClassification traitInstance = new AtlasClassification(traitName);
          traitInstance.setAttribute("type", "SSN");
  
          final String guid = createHiveTable().getGuid();
          atlasClientV2.addClassifications(guid, Collections.singletonList(traitInstance));
  
          // verify the response
          AtlasEntity withAssociationByGuid = atlasClientV2.getEntityByGuid(guid).getEntity();
          assertNotNull(withAssociationByGuid);
          assertFalse(withAssociationByGuid.getClassifications().isEmpty());
  
          boolean found = false;
          for (AtlasClassification atlasClassification : withAssociationByGuid.getClassifications()) {
              String attribute = (String) atlasClassification.getAttribute("type");
              if (attribute != null && attribute.equals("SSN")) {
                  found = true;
                  break;
              }
          }
          assertTrue(found);
      }
  
      @Test(expectedExceptions = AtlasServiceException.class)
      public void testAddTraitWithNoRegistration() throws Exception {
          final String traitName = "PII_Trait" + randomString();
          AtlasTypeUtil.createTraitTypeDef(traitName, Collections.emptySet());
  
          AtlasClassification traitInstance = new AtlasClassification(traitName);
  
          atlasClientV2.addClassifications("random", Collections.singletonList(traitInstance));
      }
  
      @Test(dependsOnMethods = "testAddTrait")
      public void testDeleteTrait() throws Exception {
          final String guid = createHiveTable().getGuid();
  
          try {
              atlasClientV2.deleteClassification(guid, traitName);
          } catch (AtlasServiceException ex) {
              fail("Deletion should've succeeded");
          }
          assertEntityAudit(guid, EntityAuditEvent.EntityAuditAction.TAG_DELETE);
      }
  
      @Test
      public void testDeleteTraitNonExistent() {
          final String traitName = "blah_trait";
  
          try {
              atlasClientV2.deleteClassification("random", traitName);
              fail("Deletion for bogus names shouldn't have succeeded");
          } catch (AtlasServiceException ex) {
              assertNotNull(ex.getStatus());
  //            assertEquals(ex.getStatus(), ClientResponse.Status.NOT_FOUND);
              assertEquals(ex.getStatus(), ClientResponse.Status.BAD_REQUEST);
              // Should it be a 400 or 404
          }
      }
  
      @Test(dependsOnMethods = "testSubmitEntity")
      public void testDeleteExistentTraitNonExistentForEntity() throws Exception {
          final String guid      = createHiveTable().getGuid();
          final String traitName = "PII_Trait" + randomString();
          AtlasClassificationDef piiTrait = AtlasTypeUtil
                  .createTraitTypeDef(traitName, Collections.emptySet(),
                          AtlasTypeUtil.createRequiredAttrDef("type", "string"));
          AtlasTypesDef typesDef = new AtlasTypesDef();
          typesDef.getClassificationDefs().add(piiTrait);
          createType(typesDef);
  
          try {
              atlasClientV2.deleteClassification(guid, traitName);
              fail("Deletion should've failed for non-existent trait association");
          } catch (AtlasServiceException ex) {
              assertNotNull(ex.getStatus());
              assertEquals(ex.getStatus(), ClientResponse.Status.BAD_REQUEST);
          }
      }
  
      @Test(dependsOnMethods = "testSubmitEntity")
      public void testGetEntityHeaderByGuid() throws Exception {
          AtlasEntityHeader header = atlasClientV2.getEntityHeaderByGuid(createHiveTable().getGuid());
          assertNotNull(header);
          assertEquals(header.getGuid(), createHiveTable().getGuid());
      }
  
      @Test(dependsOnMethods = "testSubmitEntity")
      public void testGetEntityHeaderByAttribute() throws Exception {
          AtlasEntity       hiveDB        = createHiveDB();
          String            qualifiedName = (String) hiveDB.getAttribute(NAME);
          AtlasEntityHeader header        = atlasClientV2.getEntityHeaderByAttribute(DATABASE_TYPE_V2, toMap(NAME, qualifiedName));
          assertNotNull(header);
          assertEquals(header.getTypeName(), DATABASE_TYPE_V2);
          assertEquals(header.getAttribute(NAME), qualifiedName);
      }
  
      @Test
      public void testUTF8() throws Exception {
          String classType = randomString();
          String attrName  = randomUTF8();
          String attrValue = randomUTF8();
  
          AtlasEntityDef classTypeDef = AtlasTypeUtil
                  .createClassTypeDef(classType, Collections.emptySet(),
                          AtlasTypeUtil.createUniqueRequiredAttrDef(attrName, "string"));
          AtlasTypesDef atlasTypesDef = new AtlasTypesDef();
          atlasTypesDef.getEntityDefs().add(classTypeDef);
          createType(atlasTypesDef);
  
          AtlasEntity instance = new AtlasEntity(classType);
          instance.setAttribute(attrName, attrValue);
          AtlasEntityHeader entity = createEntity(instance);
          assertNotNull(entity);
          assertNotNull(entity.getGuid());
  
          AtlasEntity entityByGuid = getEntityByGuid(entity.getGuid());
          assertEquals(entityByGuid.getAttribute(attrName), attrValue);
      }
  
      @Test(dependsOnMethods = "testSubmitEntity")
      public void testPartialUpdate() throws Exception {
          final List<AtlasEntity> columns = new ArrayList<>();
          Map<String, Object>     values  = new HashMap<>();
          values.put("name", "col1");
          values.put(NAME, "qualifiedName.col1");
          values.put("type", "string");
          values.put("comment", "col1 comment");
  
          AtlasEntity colEntity = new AtlasEntity(BaseResourceIT.COLUMN_TYPE_V2, values);
          columns.add(colEntity);
          AtlasEntity hiveTable    = createHiveTable();
          AtlasEntity tableUpdated = hiveTable;
  
          hiveTable.setAttribute("columns", AtlasTypeUtil.toObjectIds(columns));
  
          AtlasEntityWithExtInfo entityInfo = new AtlasEntityWithExtInfo(tableUpdated);
          entityInfo.addReferredEntity(colEntity);
  
          LOG.debug("Full Update entity= {}", tableUpdated);
          EntityMutationResponse updateResult = atlasClientV2.updateEntity(entityInfo);
          assertNotNull(updateResult);
          assertNotNull(updateResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
          assertFalse(updateResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE).isEmpty());
  
          String      guid          = hiveTable.getGuid();
          AtlasEntity entityByGuid1 = getEntityByGuid(guid);
          assertNotNull(entityByGuid1);
          entityByGuid1.getAttribute("columns");
  
          values.put("type", "int");
          colEntity = new AtlasEntity(BaseResourceIT.COLUMN_TYPE_V2, values);
          columns.clear();
          columns.add(colEntity);
  
          tableUpdated = new AtlasEntity(HIVE_TABLE_TYPE_V2, "name", entityByGuid1.getAttribute("name"));
          tableUpdated.setGuid(entityByGuid1.getGuid());
          tableUpdated.setAttribute("columns", AtlasTypeUtil.toObjectIds(columns));
  
          // tableUpdated = hiveTable;
          // tableUpdated.setAttribute("columns", AtlasTypeUtil.toObjectIds(columns));
  
          LOG.debug("Partial Update entity by unique attributes= {}", tableUpdated);
  
          Map<String, String> uniqAttributes = new HashMap<>();
          uniqAttributes.put(AtlasClient.REFERENCEABLE_ATTRIBUTE_NAME, (String) hiveTable.getAttribute("name"));
  
          entityInfo = new AtlasEntityWithExtInfo(tableUpdated);
          entityInfo.addReferredEntity(colEntity);
  
          EntityMutationResponse updateResponse = atlasClientV2.updateEntityByAttribute(BaseResourceIT.HIVE_TABLE_TYPE_V2, uniqAttributes, entityInfo);
  
          assertNotNull(updateResponse);
          assertNotNull(updateResponse.getEntitiesByOperation(EntityMutations.EntityOperation.PARTIAL_UPDATE));
          assertFalse(updateResponse.getEntitiesByOperation(EntityMutations.EntityOperation.PARTIAL_UPDATE).isEmpty());
  
          AtlasEntity entityByGuid2 = getEntityByGuid(guid);
          assertNotNull(entityByGuid2);
      }
  
      @Test(dependsOnMethods = "testSubmitEntity")
      public void testCompleteUpdate() throws Exception {
          final List<AtlasEntity> columns = new ArrayList<>();
          Map<String, Object>     values1 = new HashMap<>();
          values1.put("name", "col3");
          values1.put(NAME, "qualifiedName.col3");
          values1.put("type", "string");
          values1.put("comment", "col3 comment");
  
          Map<String, Object> values2 = new HashMap<>();
          values2.put("name", "col4");
          values2.put(NAME, "qualifiedName.col4");
          values2.put("type", "string");
          values2.put("comment", "col4 comment");
  
          AtlasEntity colEntity1 = new AtlasEntity(BaseResourceIT.COLUMN_TYPE_V2, values1);
          AtlasEntity colEntity2 = new AtlasEntity(BaseResourceIT.COLUMN_TYPE_V2, values2);
          columns.add(colEntity1);
          columns.add(colEntity2);
          AtlasEntity hiveTable = createHiveTable();
          hiveTable.setAttribute("columns", AtlasTypeUtil.toObjectIds(columns));
  
          AtlasEntityWithExtInfo entityInfo = new AtlasEntityWithExtInfo(hiveTable);
          entityInfo.addReferredEntity(colEntity1);
          entityInfo.addReferredEntity(colEntity2);
  
          EntityMutationResponse updateEntityResult = atlasClientV2.updateEntity(entityInfo);
          assertNotNull(updateEntityResult);
          assertNotNull(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
          assertNotNull(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE));
          //2 columns are being created, and 1 hiveTable is being updated
          assertEquals(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE).size(), 1);
          assertEquals(updateEntityResult.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE).size(), 2);
  
          AtlasEntity         entityByGuid = getEntityByGuid(hiveTable.getGuid());
          List<AtlasObjectId> refs         = (List<AtlasObjectId>) entityByGuid.getAttribute("columns");
          assertEquals(refs.size(), 2);
      }
  
      @Test
      public void testDeleteEntities() throws Exception {
          // Create 2 database entities
          AtlasEntityHeader entity1Header = createRandomDatabaseEntity();
          AtlasEntityHeader entity2Header = createRandomDatabaseEntity();
  
          // Delete the database entities
          EntityMutationResponse deleteResponse = atlasClientV2.deleteEntitiesByGuids(Arrays.asList(entity1Header.getGuid(), entity2Header.getGuid()));
  
          // Verify that deleteEntities() response has database entity guids
          assertNotNull(deleteResponse);
          assertNotNull(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE));
          assertEquals(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 2);
  
          // Verify entities were deleted from the repository.
      }
  
      @Test
      public void testPurgeEntities() throws Exception {
          // Create 2 database entities
          AtlasEntityHeader entity1Header = createRandomDatabaseEntity();
          AtlasEntityHeader entity2Header = createRandomDatabaseEntity();
  
          ApplicationProperties.get().setProperty(ENTITY_NOTIFICATION_VERSION_PROPERTY, "v2");
  
          // Delete the database entities
          EntityMutationResponse deleteResponse = atlasClientV2.deleteEntitiesByGuids(Arrays.asList(entity1Header.getGuid(), entity2Header.getGuid()));
  
          // Verify that deleteEntities() response has database entity guids
          assertNotNull(deleteResponse);
          assertNotNull(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE));
          assertEquals(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 2);
  
          //Wait for delete operation
          Thread.sleep(5000);
  
          // Purge the database entities
          Set<String>            guids         = Stream.of(entity1Header.getGuid(), entity2Header.getGuid()).collect(Collectors.toSet());
          EntityMutationResponse purgeResponse = atlasClientV2.purgeEntitiesByGuids(guids);
  
          // Verify that purgeEntities() response has database entity guids
          assertNotNull(purgeResponse);
          assertNotNull(purgeResponse.getEntitiesByOperation(EntityMutations.EntityOperation.PURGE));
          assertEquals(purgeResponse.getEntitiesByOperation(EntityMutations.EntityOperation.PURGE).size(), 2);
  
          AuditSearchParameters auditSearchParameters = TestResourceFileUtils.readObjectFromJson("audit-search-parameter-purge",
                  AuditSearchParameters.class);
          List<AtlasAuditEntry> res = atlasClientV2.getAtlasAuditByOperation(auditSearchParameters);
          // Verify that the audit entry is set
          assertNotNull(res);
      }
  
      @Test
      public void testPurgeEntitiesWithoutDelete() throws Exception {
          // Create 2 database entities
          AtlasEntityHeader entity1Header = createRandomDatabaseEntity();
          AtlasEntityHeader entity2Header = createRandomDatabaseEntity();
  
          // Purge the database entities without delete
          Set<String>            guids         = Stream.of(entity1Header.getGuid(), entity2Header.getGuid()).collect(Collectors.toSet());
          EntityMutationResponse purgeResponse = atlasClientV2.purgeEntitiesByGuids(guids);
  
          // Verify that purgeEntities() response has database entity guids
          assertNotNull(purgeResponse);
          assertNull(purgeResponse.getEntitiesByOperation(EntityMutations.EntityOperation.PURGE));
      }
  
      @Test
      public void testDeleteEntityByUniqAttribute() throws Exception {
          // Create database entity
          AtlasEntity hiveDB = createHiveDB(databaseName + randomUTF8());
  
          // Delete the database entity
          EntityMutationResponse deleteResponse = atlasClientV2.deleteEntityByAttribute(DATABASE_TYPE_V2, toMap(NAME, (String) hiveDB.getAttribute(NAME)));
  
          // Verify that deleteEntities() response has database entity guids
          assertNotNull(deleteResponse);
          assertNotNull(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE));
          assertEquals(deleteResponse.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 1);
  
          // Verify entities were deleted from the repository.
      }
  
      @Test
      public void testPartialUpdateEntityByGuid() throws Exception {
          EntityMutationResponse updateResponse = atlasClientV2.partialUpdateEntityByGuid(createHiveTable().getGuid(), Collections.singletonMap("key1", "value1"), "description");
          assertNotNull(updateResponse);
          assertNotNull(updateResponse.getEntitiesByOperation(EntityMutations.EntityOperation.PARTIAL_UPDATE));
          assertFalse(updateResponse.getEntitiesByOperation(EntityMutations.EntityOperation.PARTIAL_UPDATE).isEmpty());
      }
  
      @Test(dependsOnMethods = "testGetTraitNames")
      public void testAddClassificationsByUniqueAttribute() throws Exception {
          traitName = "PII_Trait" + randomString();
          piiTrait  =
                  AtlasTypeUtil.createTraitTypeDef(traitName, Collections.emptySet());
          AtlasTypesDef typesDef = new AtlasTypesDef();
          typesDef.getClassificationDefs().add(piiTrait);
          createType(typesDef);
  
          createdTableQualifiedName = (String) createHiveTable().getAttribute(QUALIFIED_NAME);
          atlasClientV2.addClassifications(createHiveTable().getTypeName(), toMap(QUALIFIED_NAME, createdTableQualifiedName), Collections.singletonList(new AtlasClassification(piiTrait.getName())));
          assertEntityAudit(createHiveTable().getGuid(), EntityAuditEvent.EntityAuditAction.TAG_ADD);
      }
  
      @Test(dependsOnMethods = "testAddClassificationsByUniqueAttribute")
      public void testUpdateClassifications() throws Exception {
          atlasClientV2.updateClassifications(createHiveTable().getGuid(), Collections.singletonList(new AtlasClassification(piiTrait.getName())));
          assertEntityAudit(createHiveTable().getGuid(), EntityAuditEvent.EntityAuditAction.TAG_UPDATE);
      }
  
      @Test(dependsOnMethods = "testUpdateClassifications")
      public void testUpdateClassificationsByUniqueAttribute() throws Exception {
          createdTableQualifiedName = (String) createHiveTable().getAttribute(QUALIFIED_NAME);
          atlasClientV2.updateClassifications(createHiveTable().getTypeName(), toMap(QUALIFIED_NAME, createdTableQualifiedName), Collections.singletonList(new AtlasClassification(piiTrait.getName())));
          assertEntityAudit(createHiveTable().getGuid(), EntityAuditEvent.EntityAuditAction.TAG_ADD);
      }
  
      @Test(dependsOnMethods = "testUpdateClassificationsByUniqueAttribute")
      public void testRemoveEntityClassification() throws Exception {
          createdTableQualifiedName = (String) createHiveTable().getAttribute(QUALIFIED_NAME);
          atlasClientV2.removeClassification(createHiveTable().getTypeName(), toMap(QUALIFIED_NAME, createdTableQualifiedName), piiTrait.getName());
          assertEntityAuditV2(createHiveTable().getGuid(), EntityAuditEventV2.EntityAuditActionV2.CLASSIFICATION_DELETE);
      }
  
      @Test
      public void testAddOrUpdateBusinessAttributes() throws Exception {
          Map<String, Map<String, Object>> businessAttributesMap = new HashMap<>();
          Map<String, Object>              bmAttrMapReq          = new HashMap<>();
          bmAttrMapReq.put("attr8", "01234567890123456789");
          businessAttributesMap.put("bmWithAllTypes", bmAttrMapReq);
          atlasClientV2.addOrUpdateBusinessAttributes(createHiveTable().getGuid(), false, businessAttributesMap);
          AtlasEntity.AtlasEntityWithExtInfo info = atlasClientV2.getEntityByGuid(createHiveTable().getGuid());
          assertNotNull(info);
          Map<String, Map<String, Object>> outputMap = info.getEntity().getBusinessAttributes();
          assertNotNull(outputMap);
          assertEquals(outputMap.get("bmWithAllTypes").size(), 1);
      }
  
      @Test(dependsOnMethods = "testAddOrUpdateBusinessAttributes")
      public void testRemoveBusinessAttributes() throws Exception {
          Map<String, Map<String, Object>> businessAttributesMap = new HashMap<>();
          Map<String, Object>              bmAttrMapReq          = new HashMap<>();
          bmAttrMapReq.put("attr8", "01234567890123456789");
          businessAttributesMap.put("bmWithAllTypes", bmAttrMapReq);
          atlasClientV2.removeBusinessAttributes(createHiveTable().getGuid(), businessAttributesMap);
          AtlasEntity.AtlasEntityWithExtInfo info = atlasClientV2.getEntityByGuid(createHiveTable().getGuid());
          assertNotNull(info);
          Map<String, Map<String, Object>> outputMap = info.getEntity().getBusinessAttributes();
          assertNull(outputMap);
      }
  
      // TODO Enable this test case after fixing addOrUpdateBusinessAttributesByBName API.
      @Test(enabled = false)
      public void testAddOrUpdateBusinessAttributesByBName() throws Exception {
          Map<String, Map<String, Object>> businessAttributesMap = new HashMap<>();
          Map<String, Object>              bmAttrMapReq          = new HashMap<>();
          bmAttrMapReq.put("attr8", "01234567890123456789");
          businessAttributesMap.put("bmWithAllTypes", bmAttrMapReq);
          atlasClientV2.addOrUpdateBusinessAttributes(createHiveTable().getGuid(), "bmWithAllTypes", businessAttributesMap);
          AtlasEntity.AtlasEntityWithExtInfo info = atlasClientV2.getEntityByGuid(createHiveTable().getGuid());
          assertNotNull(info);
          Map<String, Map<String, Object>> outputMap = info.getEntity().getBusinessAttributes();
          assertNotNull(outputMap);
          assertEquals(outputMap.get("bmWithAllTypes").size(), 1);
      }
  
      // TODO Enable this test case after fixing addOrUpdateBusinessAttributesByBName API.
      @Test(enabled = false, dependsOnMethods = "testAddOrUpdateBusinessAttributesByBName")
      public void testRemoveBusinessAttributesByBName() throws Exception {
          Map<String, Map<String, Object>> businessAttributesMap = new HashMap<>();
          Map<String, Object>              bmAttrMapReq          = new HashMap<>();
          bmAttrMapReq.put("attr8", "01234567890123456789");
          businessAttributesMap.put("bmWithAllTypes", bmAttrMapReq);
          atlasClientV2.removeBusinessAttributes(createHiveTable().getGuid(), "bmWithAllTypes", businessAttributesMap);
          AtlasEntity.AtlasEntityWithExtInfo info = atlasClientV2.getEntityByGuid(createHiveTable().getGuid());
          assertNotNull(info);
          Map<String, Map<String, Object>> outputMap = info.getEntity().getBusinessAttributes();
          assertNull(outputMap);
      }
  
      @Test
      public void testAddLabelsByTypeName() throws Exception {
          createdDBName = (String) createHiveDB().getAttribute(NAME);
          Set<String> labels = new HashSet<>();
          labels.add("labelByTypeName");
          atlasClientV2.setLabels(createHiveDB().getTypeName(), toMap(NAME, createdDBName), labels);
          AtlasEntityWithExtInfo info = atlasClientV2.getEntityByGuid(createHiveDB().getGuid(), false, true);
          assertNotNull(info);
          assertNotNull(info.getEntity().getLabels());
          assertEquals(info.getEntity().getLabels().size(), 1);
      }
  
      @Test(dependsOnMethods = "testAddLabelsByTypeName")
      public void testSetLabelsByTypeName() throws Exception {
          createdDBName = (String) createHiveDB().getAttribute(NAME);
          Set<String> labels = new HashSet<>();
          labels.add("labelByTypeNameNext");
          atlasClientV2.addLabels(createHiveDB().getTypeName(), toMap(NAME, createdDBName), labels);
          AtlasEntityWithExtInfo infoForSet = atlasClientV2.getEntityByGuid(createHiveDB().getGuid(), false, true);
          assertNotNull(infoForSet);
          assertNotNull(infoForSet.getEntity().getLabels());
          assertEquals(infoForSet.getEntity().getLabels().size(), 2);
      }
  
      @Test(dependsOnMethods = "testSetLabelsByTypeName")
      public void testDeleteLabelsByTypeName() throws Exception {
          Set<String> labels = new HashSet<>();
          labels.add("labelByTypeNameNext");
          createdDBName = (String) createHiveDB().getAttribute(NAME);
          atlasClientV2.removeLabels(createHiveDB().getTypeName(), toMap(NAME, createdDBName), labels);
          AtlasEntityWithExtInfo info = atlasClientV2.getEntityByGuid(createHiveDB().getGuid(), false, true);
          assertNotNull(info);
          assertNotNull(info.getEntity().getLabels());
          assertEquals(info.getEntity().getLabels().size(), 1);
      }
  
      @Test
      public void testAddClassification() throws Exception {
          traitName = "PII_Trait" + randomString();
  
          AtlasClassificationDef piiTrait = AtlasTypeUtil.createTraitTypeDef(traitName, Collections.emptySet());
  
          AtlasTypesDef typesDef = new AtlasTypesDef();
  
          typesDef.getClassificationDefs().add(piiTrait);
          createType(typesDef);
  
          ClassificationAssociateRequest request = new ClassificationAssociateRequest();
  
          request.setEntityGuids(Arrays.asList(createHiveTable().getGuid(), createHiveDB().getGuid()));
          request.setClassification(new AtlasClassification(piiTrait.getName()));
  
          atlasClientV2.addClassification(request);
  
          assertEntityAuditV2(createHiveTable().getGuid(), EntityAuditEventV2.EntityAuditActionV2.CLASSIFICATION_ADD);
          assertEntityAuditV2(createHiveDB().getGuid(), EntityAuditEventV2.EntityAuditActionV2.CLASSIFICATION_ADD);
  
          AtlasClassifications classificationsTable = atlasClientV2.getEntityClassifications(createHiveTable().getGuid(), piiTrait.getName());
  
          assertNotNull(classificationsTable);
  
          AtlasClassifications classificationsDB = atlasClientV2.getEntityClassifications(createHiveDB().getGuid(), piiTrait.getName());
  
          assertNotNull(classificationsDB);
      }
  
      @Test
      public void testDeleteClassifications() throws Exception {
          final String guid = createHiveTable().getGuid();
  
          try {
              atlasClientV2.deleteClassifications(guid, Arrays.asList(new AtlasClassification(getAndAddClassification().getName()), new AtlasClassification(getAndAddClassification().getName())));
          } catch (AtlasServiceException ex) {
              fail("Deletion should've succeeded");
          }
  
          assertEntityAudit(guid, EntityAuditEvent.EntityAuditAction.TAG_DELETE);
      }
  
      @Test
      public void testRemoveEntityClassificationByGuid() throws Exception {
          final String guid = createHiveTable().getGuid();
  
          try {
              String name = getAndAddClassification().getName();
  
              atlasClientV2.removeClassification(guid, name, guid);
          } catch (AtlasServiceException ex) {
              fail("Deletion should've succeeded");
          }
  
          assertEntityAudit(guid, EntityAuditEvent.EntityAuditAction.TAG_DELETE);
      }
  
      @Test
      public void testProduceTemplate() {
          try {
              String template = atlasClientV2.getTemplateForBulkUpdateBusinessAttributes();
  
              assertNotNull(template);
          } catch (AtlasServiceException ex) {
              fail("Deletion should've succeeded");
          }
      }
  
      //TODO Enable this test after fixing the BulkImportResponse Deserialization issue.
      @Test(dependsOnMethods = "testSubmitEntity", enabled = false)
      public void testImportBMAttributes() throws AtlasServiceException {
          BulkImportResponse response = atlasClientV2.bulkUpdateBusinessAttributes(TestResourceFileUtils.getTestFilePath("template_metadata.csv"));
  
          assertNotNull(response);
      }
  
      private String addNewType() throws Exception {
          String         typeName     = "test" + randomString();
          AtlasEntityDef classTypeDef = AtlasTypeUtil.createClassTypeDef(typeName, Collections.emptySet(), AtlasTypeUtil.createRequiredAttrDef("name", "string"), AtlasTypeUtil.createRequiredAttrDef("description", "string"));
          AtlasTypesDef  typesDef     = new AtlasTypesDef();
  
          typesDef.getEntityDefs().add(classTypeDef);
  
          createType(typesDef);
  
          return typeName;
      }
  
      private AtlasEntity getEntityByGuid(String guid) throws AtlasServiceException {
          return atlasClientV2.getEntityByGuid(guid).getEntity();
      }
  
      private void createImportFile() throws Exception {
          try {
              String filePath = TestResourceFileUtils.getTestFilePath("template_metadata.csv");
              String dbName   = (String) createHiveTable().getAttribute("name");
              String header   = "TypeName,UniqueAttributeValue,BusinessAttributeName,BusinessAttributeValue,UniqueAttributeName[optional]";
              String values   = "hive_table_v2," + dbName + ",bmWithAllTypes.attr8,\"Awesome Attribute 1\",qualifiedName";
              File   tempFile = new File(filePath);
  
              FileUtils.writeLines(tempFile, Arrays.asList(header, values));
          } catch (IOException e) {
              fail("Should have created file");
  
              throw new AtlasServiceException(e);
          }
      }
  
      private void assertEntityAudit(String dbid, EntityAuditEvent.EntityAuditAction auditAction) throws Exception {
          List<EntityAuditEvent> events = atlasClientV1.getEntityAuditEvents(dbid, (short) 100);
  
          for (EntityAuditEvent event : events) {
              if (event.getAction() == auditAction) {
                  return;
              }
          }
  
          fail("Expected audit event with action = " + auditAction);
      }
  
      private void assertEntityAuditV2(String guid, EntityAuditEventV2.EntityAuditActionV2 auditAction) throws Exception {
          // Passing auditAction as "null" as this feature is not added for InMemoryEntityRepository for IT testing.
          List<EntityAuditEventV2> events = atlasClientV2.getAuditEvents(guid, "", null, (short) 100);
  
          assertNotNull(events);
          assertNotEquals(events.size(), 0);
  
          ObjectMapper mapper = new ObjectMapper();
  
          List<EntityAuditEventV2> auditEventV2s = mapper.convertValue(events, new TypeReference<List<EntityAuditEventV2>>(){});
  
          for (EntityAuditEventV2 event : auditEventV2s) {
              if (event.getAction() == auditAction) {
                  return;
              }
          }
  
          fail("Expected audit event with action = " + auditAction);
      }
  
      private void addProperty(String guid, String property, Object value) throws AtlasServiceException {
          AtlasEntity entityByGuid = getEntityByGuid(guid);
  
          entityByGuid.setAttribute(property, value);
  
          EntityMutationResponse response = atlasClientV2.updateEntity(new AtlasEntityWithExtInfo(entityByGuid));
  
          assertNotNull(response);
          assertNotNull(response.getEntitiesByOperation(EntityMutations.EntityOperation.UPDATE));
      }
  
      private AtlasEntity createHiveDB() {
          if (dbEntity == null) {
              dbEntity = createHiveDB(databaseName);
          }
  
          return dbEntity;
      }
  
      private AtlasEntity createHiveDB(String dbName) {
          AtlasEntity       hiveDBInstanceV2 = createHiveDBInstanceV2(dbName);
          AtlasEntityHeader entityHeader     = createEntity(hiveDBInstanceV2);
  
          assertNotNull(entityHeader);
          assertNotNull(entityHeader.getGuid());
  
          hiveDBInstanceV2.setGuid(entityHeader.getGuid());
  
          return hiveDBInstanceV2;
      }
  
      private TypesUtil.Pair<AtlasEntity, AtlasEntity> createDBAndTable() throws Exception {
          AtlasEntity dbInstanceV2        = createHiveDB();
          AtlasEntity hiveTableInstanceV2 = createHiveTable();
  
          return TypesUtil.Pair.of(dbInstanceV2, hiveTableInstanceV2);
      }
  
      private AtlasEntity createHiveTable() throws Exception {
          if (tableEntity == null) {
              tableEntity = createHiveTable(createHiveDB(), tableName);
          }
  
          return tableEntity;
      }
  
      private AtlasEntity createHiveTable(AtlasEntity dbInstanceV2, String tableName) throws Exception {
          AtlasEntity       hiveTableInstanceV2 = createHiveTableInstanceV2(dbInstanceV2, tableName);
          AtlasEntityHeader createdHeader       = createEntity(hiveTableInstanceV2);
  
          assertNotNull(createdHeader);
          assertNotNull(createdHeader.getGuid());
  
          hiveTableInstanceV2.setGuid(createdHeader.getGuid());
  
          tableEntity = hiveTableInstanceV2;
  
          return hiveTableInstanceV2;
      }
  
      private AtlasClassificationDef getAndAddClassification() throws Exception {
          String                 traitNameNext = "PII_Trait" + randomString();
          AtlasClassificationDef piiTrait      = AtlasTypeUtil.createTraitTypeDef(traitNameNext, Collections.emptySet());
          AtlasTypesDef          typesDef      = new AtlasTypesDef();
  
          typesDef.getClassificationDefs().add(piiTrait);
  
          createType(typesDef);
  
          atlasClientV2.addClassifications(createHiveTable().getGuid(), Collections.singletonList(new AtlasClassification(piiTrait.getName())));
  
          assertEntityAudit(createHiveTable().getGuid(), EntityAuditEvent.EntityAuditAction.TAG_ADD);
  
          AtlasClassifications classifications = atlasClientV2.getEntityClassifications(createHiveTable().getGuid(), piiTrait.getName());
  
          assertNotNull(classifications);
  
          return piiTrait;
      }
  
      private Map<String, String> toMap(final String name, final String value) {
          return new HashMap<>(Collections.singletonMap(name, value));
      }
  
      private AtlasEntityHeader createRandomDatabaseEntity() {
          AtlasEntity db     = new AtlasEntity(DATABASE_TYPE_V2);
          String      dbName = randomString();
  
          db.setAttribute("name", dbName);
          db.setAttribute(NAME, dbName);
          db.setAttribute("clusterName", randomString());
          db.setAttribute("description", randomString());
  
          return createEntity(db);
      }
  
      private String random() {
          return RandomStringUtils.random(10);
      }
  
      private String randomUTF8() {
          String ret = random();
  
          if (!StandardCharsets.UTF_8.equals(Charset.defaultCharset())) {
              ret = new String(ret.getBytes(), StandardCharsets.UTF_8);
          }
  
          return ret;
      }
  }
  