/**
 * 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.adapters;

import org.apache.atlas.AtlasClient;
import org.apache.atlas.RequestContext;
import org.apache.atlas.TestModules;
import org.apache.atlas.TestUtilsV2;
import org.apache.atlas.exception.AtlasBaseException;
import org.apache.atlas.glossary.GlossaryService;
import org.apache.atlas.model.discovery.AtlasSearchResult;
import org.apache.atlas.model.discovery.SearchParameters;
import org.apache.atlas.model.discovery.SearchParameters.FilterCriteria;
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.AtlasEntity.AtlasEntitiesWithExtInfo;
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.ClassificationAssociateRequest;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.apache.atlas.model.instance.EntityMutations;
import org.apache.atlas.model.typedef.AtlasTypesDef;
import org.apache.atlas.repository.store.bootstrap.AtlasTypeDefStoreInitializer;
import org.apache.atlas.store.AtlasTypeDefStore;
import org.apache.atlas.type.AtlasType;
import org.apache.atlas.type.AtlasTypeRegistry;
import org.apache.atlas.type.AtlasTypeUtil;
import org.apache.atlas.web.rest.DiscoveryREST;
import org.apache.atlas.web.rest.EntityREST;
import org.apache.commons.io.FileUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Guice;
import org.testng.annotations.Test;

import javax.inject.Inject;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.apache.atlas.TestUtilsV2.CLASSIFICATION;
import static org.apache.atlas.TestUtilsV2.COLUMN_TYPE;
import static org.apache.atlas.TestUtilsV2.DATABASE_TYPE;
import static org.apache.atlas.TestUtilsV2.FETL_CLASSIFICATION;
import static org.apache.atlas.TestUtilsV2.PHI;
import static org.apache.atlas.TestUtilsV2.PII;
import static org.apache.atlas.TestUtilsV2.TABLE_TYPE;
import static org.apache.atlas.model.discovery.SearchParameters.FilterCriteria.Condition.AND;
import static org.apache.atlas.repository.Constants.CLASSIFICATION_NAMES_KEY;
import static org.apache.atlas.repository.Constants.CUSTOM_ATTRIBUTES_PROPERTY_KEY;
import static org.apache.atlas.repository.Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY;
import static org.apache.atlas.repository.Constants.STATE_PROPERTY_KEY;
import static org.apache.atlas.repository.Constants.TIMESTAMP_PROPERTY_KEY;
import static org.apache.atlas.repository.Constants.TYPE_NAME_PROPERTY_KEY;
import static org.apache.atlas.utils.TestLoadModelUtils.createTypesAsNeeded;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNotNull;
import static org.testng.Assert.assertNull;
import static org.testng.Assert.assertTrue;

@Guice(modules = TestModules.TestOnlyModule.class)
public class TestEntitiesREST {
    private static final Logger LOG = LoggerFactory.getLogger(TestEntitiesREST.class);

    @Inject
    private AtlasTypeRegistry typeRegistry;

    @Inject
    private GlossaryService glossaryService;

    @Inject
    private AtlasTypeDefStore typeStore;

    @Inject
    private DiscoveryREST discoveryREST;

    @Inject
    private EntityREST entityREST;

    private AtlasGlossary     glossary;
    private AtlasGlossaryTerm term1;
    private AtlasEntity       dbEntity;
    private AtlasEntity       tableEntity;
    private AtlasEntity       tableEntity2;
    private List<AtlasEntity> columns;
    private List<AtlasEntity> columns2;
    private SearchParameters  searchParameters;

    private final Map<String, List<String>>        createdGuids = new HashMap<>();
    private final Map<String, AtlasClassification> tagMap       = new HashMap<>();

    public static void verifyAttributes(Map<String, Object> actual, Map<String, Object> expected) {
        for (String name : actual.keySet()) {
            LOG.info("verifying attribute {} ", name);

            if (expected.get(name) != null) {
                assertEquals(actual.get(name), expected.get(name));
            }
        }
    }

    @BeforeClass
    public void setUp() throws Exception {
        AtlasTypesDef[] testTypesDefs = new AtlasTypesDef[] {TestUtilsV2.defineHiveTypes()};

        for (AtlasTypesDef typesDef : testTypesDefs) {
            AtlasTypesDef typesToCreate = AtlasTypeDefStoreInitializer.getTypesToCreate(typesDef, typeRegistry);

            if (!typesToCreate.isEmpty()) {
                typeStore.createTypesDef(typesToCreate);
            }
        }

        loadGlossaryType();

        createEntities();

        createGlossary();

        createTerms();

        initTagMap();

        registerEntities();

        addTagTo(CLASSIFICATION, TABLE_TYPE);
    }

    @AfterMethod
    public void cleanup() {
        RequestContext.clear();
    }

    @Test
    public void testGetEntities() throws Exception {
        final AtlasEntitiesWithExtInfo response = entityREST.getByGuids(createdGuids.get(DATABASE_TYPE), false, false);
        final List<AtlasEntity>        entities = response.getEntities();

        assertNotNull(entities);
        assertEquals(entities.size(), 1);
        verifyAttributes(entities);
    }

    @Test
    public void testCustomAttributesSearch() throws Exception {
        AtlasEntity         dbWithCustomAttr = new AtlasEntity(dbEntity);
        Map<String, String> customAttr       = new HashMap<>(Collections.singletonMap("key1", "value1"));

        dbWithCustomAttr.setCustomAttributes(customAttr);

        AtlasEntitiesWithExtInfo atlasEntitiesWithExtInfo = new AtlasEntitiesWithExtInfo(dbWithCustomAttr);
        EntityMutationResponse   response                 = entityREST.createOrUpdate(atlasEntitiesWithExtInfo);

        assertNotNull(response.getUpdatedEntitiesByTypeName(DATABASE_TYPE));

        searchParameters = new SearchParameters();

        searchParameters.setTypeName("_ALL_ENTITY_TYPES");

        SearchParameters.FilterCriteria filter = new SearchParameters.FilterCriteria();

        filter.setAttributeName(CUSTOM_ATTRIBUTES_PROPERTY_KEY);
        filter.setOperator(SearchParameters.Operator.CONTAINS);
        filter.setAttributeValue("key1=value1");

        searchParameters.setEntityFilters(filter);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 1);
    }

    @Test
    public void testBasicSearch() throws Exception {
        // search entities with classification named classification
        searchParameters = new SearchParameters();

        searchParameters.setIncludeSubClassifications(true);
        searchParameters.setClassification(TestUtilsV2.CLASSIFICATION);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);
    }

    @Test
    public void testSearchByMultiSystemAttributes() throws Exception {
        searchParameters = new SearchParameters();

        searchParameters.setTypeName("_ALL_ENTITY_TYPES");

        SearchParameters.FilterCriteria fc     = new SearchParameters.FilterCriteria();
        SearchParameters.FilterCriteria subFc1 = new SearchParameters.FilterCriteria();
        SearchParameters.FilterCriteria subFc2 = new SearchParameters.FilterCriteria();

        subFc1.setAttributeName(MODIFICATION_TIMESTAMP_PROPERTY_KEY);
        subFc1.setOperator(SearchParameters.Operator.LT);
        subFc1.setAttributeValue(String.valueOf(System.currentTimeMillis()));

        subFc2.setAttributeName(TIMESTAMP_PROPERTY_KEY);
        subFc2.setOperator(SearchParameters.Operator.LT);
        subFc2.setAttributeValue(String.valueOf(System.currentTimeMillis()));

        fc.setCriterion(Arrays.asList(subFc1, subFc2));
        fc.setCondition(AND);
        searchParameters.setEntityFilters(fc);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertTrue(res.getEntities().size() > 5);
    }

    @Test
    public void testWildCardBasicSearch() throws Exception {
        //table - classification
        searchParameters = new SearchParameters();

        searchParameters.setClassification("*");

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        searchParameters.setClassification("_CLASSIFIED");

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        // Test wildcard usage of basic search
        searchParameters.setClassification("cl*");

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        searchParameters.setClassification("*ion");

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        searchParameters.setClassification("*l*");

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        searchParameters.setClassification("_NOT_CLASSIFIED");
        searchParameters.setTypeName(DATABASE_TYPE);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 1);
    }

    @Test(dependsOnMethods = "testWildCardBasicSearch")
    public void testBasicSearchWithAttr() throws Exception {
        searchParameters = new SearchParameters();

        searchParameters.setClassification("cla*");
        searchParameters.setTypeName(TABLE_TYPE);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        // table - classification
        // column - phi
        addTagTo(PHI, COLUMN_TYPE);

        FilterCriteria filterCriteria = new FilterCriteria();
        filterCriteria.setAttributeName("stringAttr");
        filterCriteria.setOperator(SearchParameters.Operator.CONTAINS);
        filterCriteria.setAttributeValue("sample");

        // basic search with tag filterCriteria
        searchParameters = new SearchParameters();

        searchParameters.setClassification(PHI);
        searchParameters.setTagFilters(filterCriteria);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        filterCriteria.setAttributeName("stringAttr");
        filterCriteria.setOperator(SearchParameters.Operator.EQ);
        filterCriteria.setAttributeValue("sample_string");

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        filterCriteria.setAttributeValue("SAMPLE_STRING");

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNull(res.getEntities());
    }

    @Test(dependsOnMethods = "testBasicSearchWithAttr")
    public void testBasicSearchWithSubTypes() throws Exception {
        // table - classification
        // column - phi
        searchParameters = new SearchParameters();

        searchParameters.setClassification(TestUtilsV2.CLASSIFICATION);
        searchParameters.setIncludeSubClassifications(true);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        // table - classification
        // database - fetl_classification
        // column - phi
        addTagTo(FETL_CLASSIFICATION, DATABASE_TYPE);

        final AtlasClassification resultTag = entityREST.getClassification(createdGuids.get(DATABASE_TYPE).get(0), TestUtilsV2.FETL_CLASSIFICATION);

        assertNotNull(resultTag);
        assertEquals(resultTag.getTypeName(), FETL_CLASSIFICATION);

        // basic search with subtypes
        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 3);

        // basic search without subtypes
        searchParameters.setIncludeSubClassifications(false);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);
    }

    @Test(dependsOnMethods = "testBasicSearchWithSubTypes")
    public void testGraphQueryFilter() throws Exception {
        // table - classification
        // database - fetl_classification
        // column - phi
        searchParameters = new SearchParameters();

        searchParameters.setQuery("sample_string");
        searchParameters.setClassification(PHI);

        SearchParameters.FilterCriteria fc = new SearchParameters.FilterCriteria();
        fc.setOperator(SearchParameters.Operator.EQ);
        fc.setAttributeName("booleanAttr");
        fc.setAttributeValue("true");

        searchParameters.setTagFilters(fc);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);
        assertEquals(res.getEntities().get(0).getTypeName(), COLUMN_TYPE);

        Map<String, Object> attributes = new HashMap<>();

        attributes.put("stringAttr", "sample_string");
        attributes.put("booleanAttr", false);
        attributes.put("integerAttr", 100);

        AtlasClassification cls = new AtlasClassification(TestUtilsV2.PHI, attributes);

        ClassificationAssociateRequest clsAssRequest = new ClassificationAssociateRequest(Collections.singletonList(createdGuids.get(TABLE_TYPE).get(0)), cls);

        entityREST.addClassification(clsAssRequest);

        final AtlasClassification resultTag = entityREST.getClassification(createdGuids.get(TABLE_TYPE).get(0), TestUtilsV2.PHI);

        assertNotNull(resultTag);
        assertEquals(resultTag.getTypeName(), PHI);

        // table - phi, classification
        // database - fetl_classification
        // column - phi
        fc.setAttributeValue("false");

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 1);
        assertEquals(res.getEntities().get(0).getTypeName(), TABLE_TYPE);
    }

    @Test(dependsOnMethods = "testGraphQueryFilter")
    public void testSearchByMultiTags() throws Exception {
        addTagTo(PHI, DATABASE_TYPE);

        // database - phi, felt_classification
        // table1 - phi, classification, table2 - classification,
        // column - phi
        searchParameters = new SearchParameters();

        searchParameters.setIncludeSubClassifications(false);
        searchParameters.setTypeName("_ALL_ENTITY_TYPES");

        SearchParameters.FilterCriteria fc     = new SearchParameters.FilterCriteria();
        SearchParameters.FilterCriteria subFc1 = new SearchParameters.FilterCriteria();
        SearchParameters.FilterCriteria subFc2 = new SearchParameters.FilterCriteria();

        subFc1.setAttributeName(CLASSIFICATION_NAMES_KEY);
        subFc1.setOperator(SearchParameters.Operator.CONTAINS);
        subFc1.setAttributeValue(PHI);

        subFc2.setAttributeName(CLASSIFICATION_NAMES_KEY);
        subFc2.setOperator(SearchParameters.Operator.CONTAINS);
        subFc2.setAttributeValue(CLASSIFICATION);

        fc.setCriterion(Arrays.asList(subFc1, subFc2));
        fc.setCondition(AND);

        searchParameters.setEntityFilters(fc);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        subFc2.setAttributeValue(FETL_CLASSIFICATION);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 1);
    }

    @Test(dependsOnMethods = "testSearchByMultiTags")
    public void testSearchByTerms() throws Exception {
        // database - phi, felt_classification
        // table1 - phi, classification, term: term1 | table2 - classification, term:term2
        // column - phi
        assignTermTo(term1, tableEntity);
        assignTermTo(term1, tableEntity2);

        searchParameters = new SearchParameters();

        searchParameters.setTermName(term1.getName() + "@testSearchGlossary");
        searchParameters.setClassification(CLASSIFICATION);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        searchParameters.setClassification(PII);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNull(res.getEntities());

        searchParameters.setClassification(PHI);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 1);
    }

    @Test(dependsOnMethods = "testSearchByMultiTags")
    public void testSearchByOtherSystemAttributes() throws Exception {
        // database - phi, felt_classification
        // table1 - phi, classification, table2 - classification,
        // col1 - phi,  col2 - phi
        searchParameters = new SearchParameters();

        searchParameters.setTypeName("_ALL_ENTITY_TYPES");
        SearchParameters.FilterCriteria fc = new SearchParameters.FilterCriteria();

        fc.setAttributeName(CLASSIFICATION_NAMES_KEY);
        fc.setOperator(SearchParameters.Operator.EQ);
        fc.setAttributeValue(CLASSIFICATION);

        searchParameters.setEntityFilters(fc);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        fc.setOperator(SearchParameters.Operator.CONTAINS);
        fc.setAttributeValue("cla");

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 3);

        fc.setOperator(SearchParameters.Operator.CONTAINS);
        fc.setAttributeValue(PHI);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 4);
    }

    @Test(dependsOnMethods = "testSearchByOtherSystemAttributes")
    public void testBasicSearchWithFilter() throws Exception {
        // database - phi, felt_classification
        // table1 - phi, classification, table2 - classification,
        // col - phi
        searchParameters = new SearchParameters();

        searchParameters.setIncludeSubClassifications(false);
        searchParameters.setClassification(TestUtilsV2.CLASSIFICATION);

        SearchParameters.FilterCriteria fc = new SearchParameters.FilterCriteria();

        fc.setOperator(SearchParameters.Operator.CONTAINS);
        fc.setAttributeValue("new comments");
        fc.setAttributeName("tag");

        searchParameters.setTagFilters(fc);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNull(res.getEntities());

        fc.setOperator(SearchParameters.Operator.ENDS_WITH);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNull(res.getEntities());

        fc.setOperator(SearchParameters.Operator.STARTS_WITH);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNull(res.getEntities());
    }

    @Test(dependsOnMethods = "testBasicSearchWithFilter")
    public void testSearchBySingleSystemAttribute() throws Exception {
        // database - phi, felt_classification
        // table1 - phi, classification, table2 - classification,
        // col - phi
        searchParameters = new SearchParameters();

        searchParameters.setTypeName("_ALL_ENTITY_TYPES");

        SearchParameters.FilterCriteria fc = new SearchParameters.FilterCriteria();

        fc.setAttributeName(STATE_PROPERTY_KEY);
        fc.setOperator(SearchParameters.Operator.EQ);
        fc.setAttributeValue("DELETED");

        searchParameters.setEntityFilters(fc);

        entityREST.deleteByGuid(createdGuids.get(DATABASE_TYPE).get(0));

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 1);

        searchParameters.setTypeName("_ALL_ENTITY_TYPES");
        fc.setAttributeName(TYPE_NAME_PROPERTY_KEY);
        fc.setOperator(SearchParameters.Operator.EQ);
        fc.setAttributeValue(TABLE_TYPE);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 2);

        searchParameters = new SearchParameters();

        searchParameters.setTypeName(TABLE_TYPE);
        fc.setAttributeName(STATE_PROPERTY_KEY);
        fc.setOperator(SearchParameters.Operator.EQ);
        fc.setAttributeValue("DELETED");

        searchParameters.setEntityFilters(fc);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNull(res.getEntities());
    }

    @Test(dependsOnMethods = "testSearchBySingleSystemAttribute")
    public void testSearchBySystemAttributesWithQuery() throws Exception {
        // database - phi, felt_classification
        // table1 - phi, classification, table2 - classification,
        // col - phi
        searchParameters = new SearchParameters();

        searchParameters.setTypeName("_ALL_ENTITY_TYPES");

        SearchParameters.FilterCriteria fc     = new SearchParameters.FilterCriteria();
        SearchParameters.FilterCriteria subFc1 = new SearchParameters.FilterCriteria();
        SearchParameters.FilterCriteria subFc2 = new SearchParameters.FilterCriteria();

        subFc1.setAttributeName(TIMESTAMP_PROPERTY_KEY);
        subFc1.setOperator(SearchParameters.Operator.LT);
        subFc1.setAttributeValue(String.valueOf(System.currentTimeMillis()));

        subFc2.setAttributeName(STATE_PROPERTY_KEY);
        subFc2.setOperator(SearchParameters.Operator.EQ);
        subFc2.setAttributeValue("DELETED");

        fc.setCriterion(Arrays.asList(subFc1, subFc2));
        fc.setCondition(AND);

        searchParameters.setEntityFilters(fc);
        searchParameters.setQuery("sample_string");

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 1);

        searchParameters = new SearchParameters();

        searchParameters.setQuery("classification");
        searchParameters.setClassification(PHI);

        res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 1);
    }

    @Test(dependsOnMethods = "testSearchBySystemAttributesWithQuery")
    public void testTagSearchBySystemAttributes() throws Exception {
        // database - phi, felt_classification
        // table1 - phi, classification, table2 - classification,
        // col - phi
        searchParameters = new SearchParameters();

        searchParameters.setClassification("_ALL_CLASSIFICATION_TYPES");

        SearchParameters.FilterCriteria fc     = new SearchParameters.FilterCriteria();
        SearchParameters.FilterCriteria subFc1 = new SearchParameters.FilterCriteria();
        SearchParameters.FilterCriteria subFc2 = new SearchParameters.FilterCriteria();

        subFc1.setAttributeName(TIMESTAMP_PROPERTY_KEY);
        subFc1.setOperator(SearchParameters.Operator.LT);
        subFc1.setAttributeValue(String.valueOf(System.currentTimeMillis()));

        subFc2.setAttributeName(STATE_PROPERTY_KEY);
        subFc2.setOperator(SearchParameters.Operator.EQ);
        subFc2.setAttributeValue("ACTIVE");

        fc.setCriterion(Arrays.asList(subFc1, subFc2));
        fc.setCondition(AND);

        searchParameters.setTagFilters(fc);

        AtlasSearchResult res = discoveryREST.searchWithParameters(searchParameters);

        assertNotNull(res.getEntities());
        assertEquals(res.getEntities().size(), 7);
    }

    /* Disabled until EntityREST.deleteByIds() is implemented
     *
    @Test(dependsOnMethods = "testGetEntities")
    public void testDeleteEntities() throws Exception {
        final EntityMutationResponse response = entityREST.deleteByGuids(createdGuids);
        final List<AtlasEntityHeader> entities = response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE);

        assertNotNull(entities);
        assertEquals(entities.size(), 3);
    }
    *
    */

    @Test(dependsOnMethods = "testTagSearchBySystemAttributes")
    public void testUpdateWithSerializedEntities() throws Exception {
        //Check with serialization and deserialization of entity attributes for the case
        // where attributes which are de-serialized into a map
        AtlasEntity dbEntity    = TestUtilsV2.createDBEntity();
        AtlasEntity tableEntity = TestUtilsV2.createTableEntity(dbEntity);

        final AtlasEntity colEntity = TestUtilsV2.createColumnEntity(tableEntity);
        List<AtlasEntity> columns   = new ArrayList<>(Collections.singletonList(colEntity));

        tableEntity.setAttribute("columns", getObjIdList(columns));

        AtlasEntity newDBEntity    = serDeserEntity(dbEntity);
        AtlasEntity newTableEntity = serDeserEntity(tableEntity);

        AtlasEntitiesWithExtInfo newEntities = new AtlasEntitiesWithExtInfo();

        newEntities.addEntity(newDBEntity);
        newEntities.addEntity(newTableEntity);

        for (AtlasEntity column : columns) {
            newEntities.addReferredEntity(serDeserEntity(column));
        }

        EntityMutationResponse response2 = entityREST.createOrUpdate(newEntities);

        List<AtlasEntityHeader> newGuids = response2.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);

        assertNotNull(newGuids);
        assertEquals(newGuids.size(), 3);
    }

    AtlasEntity serDeserEntity(AtlasEntity entity) {
        //Convert from json to object and back to trigger the case where it gets translated to a map for attributes instead of AtlasEntity
        String      jsonString = AtlasType.toJson(entity);

        return AtlasType.fromJson(jsonString, AtlasEntity.class);
    }

    private void loadGlossaryType() throws IOException, AtlasBaseException {
        registerAtlasTypesDef("/addons/models/0000-Area0/0010-base_model.json");
        registerAtlasTypesDef("/addons/models/0000-Area0/0011-glossary_model.json");
    }

    private void registerAtlasTypesDef(String path) throws IOException, AtlasBaseException {
        String projectBaseDirectory = System.getProperty("projectBaseDir");

        String baseModel = projectBaseDirectory + path;
        File   f         = new File(baseModel);
        String s         = FileUtils.readFileToString(f);

        createTypesAsNeeded(AtlasType.fromJson(s, AtlasTypesDef.class), typeStore, typeRegistry);
    }

    private void createGlossary() throws AtlasBaseException {
        glossary = new AtlasGlossary();

        glossary.setQualifiedName("testSearchGlossary");
        glossary.setName("Search glossary");
        glossary.setShortDescription("Short description");
        glossary.setLongDescription("Long description");
        glossary.setUsage("N/A");
        glossary.setLanguage("en-US");

        AtlasGlossary created = glossaryService.createGlossary(glossary);

        glossary.setGuid(created.getGuid());
    }

    private void assignTermTo(AtlasGlossaryTerm term, AtlasEntity entity) throws AtlasBaseException {
        AtlasRelatedObjectId relatedObjectId = new AtlasRelatedObjectId();

        relatedObjectId.setGuid(entity.getGuid());
        relatedObjectId.setTypeName(entity.getTypeName());

        glossaryService.assignTermToEntities(term.getGuid(), Collections.singletonList(relatedObjectId));
    }

    private void createTerms() throws AtlasBaseException {
        term1 = new AtlasGlossaryTerm();

        // Glossary anchor
        AtlasGlossaryHeader glossaryId = new AtlasGlossaryHeader();

        glossaryId.setGlossaryGuid(glossary.getGuid());

        term1.setName("term1");
        term1.setShortDescription("Short description");
        term1.setLongDescription("Long description");
        term1.setAbbreviation("CHK");
        term1.setExamples(Arrays.asList("Personal", "Joint"));
        term1.setUsage("N/A");
        term1.setAnchor(glossaryId);

        AtlasGlossaryTerm created1 = glossaryService.createTerm(term1);

        term1.setGuid(created1.getGuid());
    }

    private void createEntities() {
        dbEntity     = TestUtilsV2.createDBEntity();
        tableEntity  = TestUtilsV2.createTableEntity(dbEntity);
        tableEntity2 = TestUtilsV2.createTableEntity(dbEntity);

        final AtlasEntity colEntity  = TestUtilsV2.createColumnEntity(tableEntity);
        final AtlasEntity colEntity2 = TestUtilsV2.createColumnEntity(tableEntity2);

        columns  = new ArrayList<>(Collections.singletonList(colEntity));
        columns2 = new ArrayList<>(Collections.singletonList(colEntity2));

        tableEntity.setAttribute("columns", getObjIdList(columns));
        tableEntity2.setAttribute("columns", getObjIdList(columns2));
    }

    private void initTagMap() {
        Map<String, Object> attributes = new HashMap<>();

        attributes.put("stringAttr", "sample_string");
        attributes.put("booleanAttr", true);
        attributes.put("integerAttr", 100);

        tagMap.put(PHI, new AtlasClassification(TestUtilsV2.PHI, attributes));

        tagMap.put(PII, new AtlasClassification(TestUtilsV2.PII, new HashMap<>(Collections.singletonMap("cls", "clsName"))));

        tagMap.put(CLASSIFICATION, new AtlasClassification(TestUtilsV2.CLASSIFICATION, new HashMap<>(Collections.singletonMap("tag", "tagName"))));

        attributes = new HashMap<>();

        attributes.put("cls", "clsName");
        attributes.put("tag", "tagName");

        tagMap.put(FETL_CLASSIFICATION, new AtlasClassification(FETL_CLASSIFICATION, attributes));
    }

    private void addTagTo(String tagName, String type) throws Exception {
        AtlasClassification            tag                            = tagMap.get(tagName);
        ClassificationAssociateRequest classificationAssociateRequest = new ClassificationAssociateRequest(createdGuids.get(type), tag);

        entityREST.addClassification(classificationAssociateRequest);

        for (int i = 0; i < createdGuids.get(type).size() - 1; i++) {
            final AtlasClassification resultTag = entityREST.getClassification(createdGuids.get(type).get(i), tagName);

            assertNotNull(resultTag);
            assertEquals(resultTag.getTypeName(), tag.getTypeName());
        }
    }

    private void registerEntities() throws Exception {
        AtlasEntitiesWithExtInfo entities = new AtlasEntitiesWithExtInfo();

        entities.addEntity(dbEntity);
        entities.addEntity(tableEntity);
        entities.addEntity(tableEntity2);

        for (AtlasEntity column : columns) {
            entities.addReferredEntity(column);
        }

        for (AtlasEntity column : columns2) {
            entities.addReferredEntity(column);
        }

        EntityMutationResponse  response = entityREST.createOrUpdate(entities);
        List<AtlasEntityHeader> guids    = response.getEntitiesByOperation(EntityMutations.EntityOperation.CREATE);

        assertNotNull(guids);
        assertEquals(guids.size(), 5);

        for (AtlasEntityHeader header : guids) {
            if (!createdGuids.containsKey(header.getTypeName())) {
                createdGuids.put(header.getTypeName(), new ArrayList<>());
            }

            createdGuids.get(header.getTypeName()).add(header.getGuid());
        }
    }

    private void verifyAttributes(List<AtlasEntity> retrievedEntities) {
        AtlasEntity retrievedDBEntity     = null;
        AtlasEntity retrievedTableEntity  = null;
        AtlasEntity retrievedColumnEntity = null;

        for (AtlasEntity entity : retrievedEntities) {
            if (entity.getTypeName().equals(DATABASE_TYPE)) {
                retrievedDBEntity = entity;
            }

            if (entity.getTypeName().equals(TABLE_TYPE)) {
                retrievedTableEntity = entity;
            }

            if (entity.getTypeName().equals(TestUtilsV2.COLUMN_TYPE)) {
                retrievedColumnEntity = entity;
            }
        }

        if (retrievedDBEntity != null) {
            LOG.info("verifying entity of type {} ", dbEntity.getTypeName());

            verifyAttributes(retrievedDBEntity.getAttributes(), dbEntity.getAttributes());
        }

        if (retrievedColumnEntity != null) {
            LOG.info("verifying entity of type {} ", columns.get(0).getTypeName());

            assertEquals(columns.get(0).getAttribute(AtlasClient.NAME), retrievedColumnEntity.getAttribute(AtlasClient.NAME));
            assertEquals(columns.get(0).getAttribute("type"), retrievedColumnEntity.getAttribute("type"));
        }

        if (retrievedTableEntity != null) {
            LOG.info("verifying entity of type {} ", tableEntity.getTypeName());

            //String
            assertEquals(tableEntity.getAttribute(AtlasClient.NAME), retrievedTableEntity.getAttribute(AtlasClient.NAME));
            //Map
            assertEquals(tableEntity.getAttribute("parametersMap"), retrievedTableEntity.getAttribute("parametersMap"));
            //enum
            assertEquals(tableEntity.getAttribute("tableType"), retrievedTableEntity.getAttribute("tableType"));
            //date
            assertEquals(tableEntity.getAttribute("created"), retrievedTableEntity.getAttribute("created"));
            //array of Ids
            assertEquals(((List<AtlasObjectId>) retrievedTableEntity.getAttribute("columns")).get(0).getGuid(), retrievedColumnEntity.getGuid());
            //array of structs
            assertEquals(retrievedTableEntity.getAttribute("partitions"), tableEntity.getAttribute("partitions"));
        }
    }

    private List<AtlasObjectId> getObjIdList(Collection<AtlasEntity> entities) {
        List<AtlasObjectId> ret = new ArrayList<>();

        for (AtlasEntity entity : entities) {
            ret.add(AtlasTypeUtil.getAtlasObjectId(entity));
        }

        return ret;
    }
}
