/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.atlas.repository.store.graph.v2;

import org.apache.atlas.GraphTransactionInterceptor;
import org.apache.atlas.TestModules;
import org.apache.atlas.TestUtilsV2;
import org.apache.atlas.model.instance.AtlasEntity;
import org.apache.atlas.model.instance.AtlasEntity.AtlasEntitiesWithExtInfo;
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.AtlasStruct;
import org.apache.atlas.model.instance.EntityMutationResponse;
import org.apache.atlas.model.typedef.AtlasTypesDef;
import org.apache.atlas.repository.graphdb.AtlasEdge;
import org.apache.atlas.repository.graphdb.AtlasEdgeDirection;
import org.apache.atlas.repository.graphdb.AtlasVertex;
import org.apache.atlas.type.AtlasEntityType;
import org.apache.commons.lang3.time.DateUtils;
import org.locationtech.jts.util.Assert;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Guice;
import org.testng.annotations.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import static org.apache.atlas.TestUtilsV2.ENTITY_TYPE;
import static org.apache.atlas.TestUtilsV2.ENTITY_TYPE_MAP;
import static org.apache.atlas.TestUtilsV2.ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR;
import static org.apache.atlas.TestUtilsV2.ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR_DELETE;
import static org.apache.atlas.TestUtilsV2.ENTITY_TYPE_WITH_SIMPLE_ATTR;
import static org.apache.atlas.TestUtilsV2.NAME;
import static org.apache.atlas.repository.graph.GraphHelper.getStatus;
import static org.apache.atlas.type.AtlasTypeUtil.getAtlasObjectId;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNull;

@Guice(modules = TestModules.TestOnlyModule.class)
public class AtlasComplexAttributesTest extends AtlasEntityTestBase {
    private AtlasEntityWithExtInfo complexCollectionAttrEntity;
    private AtlasEntityWithExtInfo complexCollectionAttrEntityForDelete;
    private AtlasEntityWithExtInfo mapAttributesEntity;

    @BeforeClass
    public void setUp() throws Exception {
        super.setUp();

        // create typeDefs
        AtlasTypesDef[] testTypesDefs = new AtlasTypesDef[] {TestUtilsV2.defineTypeWithComplexCollectionAttributes(), TestUtilsV2.defineTypeWithMapAttributes(), TestUtilsV2.defineSimpleAttrType()};
        createTypesDef(testTypesDefs);

        // create entity
        complexCollectionAttrEntity          = TestUtilsV2.createComplexCollectionAttrEntity();
        complexCollectionAttrEntityForDelete = TestUtilsV2.createComplexCollectionAttrEntity();
        mapAttributesEntity                  = TestUtilsV2.createMapAttrEntity();
    }

    @Test
    public void testCreateComplexAttributeEntity() throws Exception {
        init();

        EntityMutationResponse response      = entityStore.createOrUpdate(new AtlasEntityStream(complexCollectionAttrEntity), false);
        AtlasEntityHeader      entityCreated = response.getFirstCreatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexCollectionAttrEntity, getEntityFromStore(entityCreated));
    }

    @Test
    public void testPrimitiveMapAttributes() throws Exception {
        init();

        EntityMutationResponse response        = entityStore.createOrUpdate(new AtlasEntityStream(mapAttributesEntity), false);
        AtlasEntityHeader      entityCreated   = response.getFirstCreatedEntityByTypeName(ENTITY_TYPE_MAP);
        AtlasEntity            entityFromStore = getEntityFromStore(entityCreated);

        validateEntity(mapAttributesEntity, entityFromStore);

        // Modify map of primitives
        AtlasEntity attrEntity = getEntityFromStore(mapAttributesEntity.getEntity().getGuid());

        Map<String, String> map1 = new HashMap<>();

        map1.put("map1Key11", "value11");
        map1.put("map1Key22", "value22");
        map1.put("map1Key33", "value33");

        Map<String, Integer> map2 = new HashMap<>();

        map2.put("map2Key11", 1100);
        map2.put("map2Key22", 2200);
        map2.put("map2Key33", 3300);

        Map<String, Boolean> map3 = new HashMap<>();

        map3.put("map3Key11", true);
        map3.put("map3Key22", false);
        map3.put("map3Key33", true);

        Map<String, Float> map4 = new HashMap<>();

        map4.put("map4Key11", 11.0f);
        map4.put("map4Key22", 22.0f);
        map4.put("map4Key33", 33.0f);

        Map<String, Date> map5 = new HashMap<>();

        map5.put("map5Key11", DateUtils.addHours(new Date(), 1));
        map5.put("map5Key22", DateUtils.addHours(new Date(), 2));
        map5.put("map5Key33", DateUtils.addHours(new Date(), 3));

        updateEntityMapAttributes(attrEntity, map1, map2, map3, map4, map5);

        AtlasEntitiesWithExtInfo attrEntitiesInfo = new AtlasEntitiesWithExtInfo(attrEntity);

        response = entityStore.createOrUpdate(new AtlasEntityStream(attrEntitiesInfo), false);

        AtlasEntityHeader updatedAttrEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_MAP);
        AtlasEntity       updatedFromStore  = getEntityFromStore(updatedAttrEntity);

        validateEntity(attrEntitiesInfo, updatedFromStore);

        // Add new entry to map of primitives
        map1.put("map1Key44", "value44");
        map2.put("map2Key44", 4400);
        map3.put("map3Key44", false);
        map4.put("map4Key44", 44.0f);
        map5.put("map5Key44", DateUtils.addHours(new Date(), 4));

        updateEntityMapAttributes(attrEntity, map1, map2, map3, map4, map5);

        attrEntitiesInfo  = new AtlasEntitiesWithExtInfo(attrEntity);
        response          = entityStore.createOrUpdate(new AtlasEntityStream(attrEntitiesInfo), false);
        updatedAttrEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_MAP);
        updatedFromStore  = getEntityFromStore(updatedAttrEntity);

        validateEntity(attrEntitiesInfo, updatedFromStore);

        // Remove an entry from map of primitives
        map1.remove("map1Key11");
        map2.remove("map2Key11");
        map3.remove("map3Key11");
        map4.remove("map4Key11");
        map5.remove("map5Key11");

        updateEntityMapAttributes(attrEntity, map1, map2, map3, map4, map5);

        attrEntitiesInfo  = new AtlasEntitiesWithExtInfo(attrEntity);
        response          = entityStore.createOrUpdate(new AtlasEntityStream(attrEntitiesInfo), false);
        updatedAttrEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_MAP);
        updatedFromStore  = getEntityFromStore(updatedAttrEntity);

        validateEntity(attrEntitiesInfo, updatedFromStore);

        // Edit existing entry to map of primitives
        map1.put("map1Key44", "value44-edit");
        map2.put("map2Key44", 5555);
        map3.put("map3Key44", true);
        map4.put("map4Key44", 55.5f);
        map5.put("map5Key44", DateUtils.addHours(new Date(), 5));

        updateEntityMapAttributes(attrEntity, map1, map2, map3, map4, map5);

        attrEntitiesInfo  = new AtlasEntitiesWithExtInfo(attrEntity);
        response          = entityStore.createOrUpdate(new AtlasEntityStream(attrEntitiesInfo), false);
        updatedAttrEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_MAP);
        updatedFromStore  = getEntityFromStore(updatedAttrEntity);

        validateEntity(attrEntitiesInfo, updatedFromStore);

        // clear primitive map entries
        map1.clear();
        map2.clear();
        map3.clear();
        map4.clear();
        map5.clear();

        updateEntityMapAttributes(attrEntity, map1, map2, map3, map4, map5);

        attrEntitiesInfo  = new AtlasEntitiesWithExtInfo(attrEntity);
        response          = entityStore.createOrUpdate(new AtlasEntityStream(attrEntitiesInfo), false);
        updatedAttrEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_MAP);
        updatedFromStore  = getEntityFromStore(updatedAttrEntity);

        validateEntity(attrEntitiesInfo, updatedFromStore);
    }

    @Test
    public void testArrayAttribute() throws Exception {
        init();

        AtlasEntityWithExtInfo simpleEntity        = TestUtilsV2.createSimpleAttrTypeEntity();
        EntityMutationResponse response            = entityStore.createOrUpdate(new AtlasEntityStream(simpleEntity), false);
        AtlasEntityHeader      simpleEntityHeader  = response.getFirstCreatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
        AtlasEntity            createdSimpleEntity = getEntityFromStore(simpleEntityHeader);

        validateEntity(simpleEntity, createdSimpleEntity);

        createdSimpleEntity.setAttribute("stringAtrr", null);
        createdSimpleEntity.setAttribute("mapOfStrings", Collections.emptyMap());
        createdSimpleEntity.setAttribute("arrayOfStrings", Collections.emptyList());
        createdSimpleEntity.setAttribute("puArray", Collections.emptyList());
        createdSimpleEntity.setAttribute("puMap", Collections.emptyMap());

        EntityMutationResponse responseUpdated           = entityStore.createOrUpdate(new AtlasEntityStream(createdSimpleEntity), false);
        AtlasEntityHeader      simpleEntityUpdatedHeader = responseUpdated.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
        AtlasEntity            updatedSimpleEntity       = getEntityFromStore(simpleEntityUpdatedHeader);

        assertNull(updatedSimpleEntity.getAttribute("stringAtrr"));
        assertEquals(updatedSimpleEntity.getAttribute("mapOfStrings"), Collections.emptyMap());
        assertEquals(updatedSimpleEntity.getAttribute("arrayOfStrings"), Collections.emptyList());
        assertEquals(updatedSimpleEntity.getAttribute("puArray"), Collections.emptyList());
        assertEquals(updatedSimpleEntity.getAttribute("puMap"), Collections.emptyMap());

        updatedSimpleEntity.setAttribute("stringAtrr", "");
        updatedSimpleEntity.setAttribute("mapOfStrings", null);
        updatedSimpleEntity.setAttribute("arrayOfStrings", null);
        updatedSimpleEntity.setAttribute("puArray", null);
        updatedSimpleEntity.setAttribute("puMap", null);

        EntityMutationResponse responseUpdatedAgain           = entityStore.createOrUpdate(new AtlasEntityStream(updatedSimpleEntity), false);
        AtlasEntityHeader      simpleEntityUpdatedAgainHeader = responseUpdatedAgain.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
        AtlasEntity            updatedAgainSimpleEntity       = getEntityFromStore(simpleEntityUpdatedAgainHeader);

        assertEquals(updatedAgainSimpleEntity.getAttribute("stringAtrr"), "");
        assertNull(updatedAgainSimpleEntity.getAttribute("arrayOfStrings"));
        assertNull(updatedAgainSimpleEntity.getAttribute("mapOfStrings"));
        assertNull(updatedAgainSimpleEntity.getAttribute("puArray"));
        assertNull(updatedAgainSimpleEntity.getAttribute("puMap"));

        updatedAgainSimpleEntity.setAttribute("stringAtrr", "Dummy String Test 3");
        updatedAgainSimpleEntity.setAttribute("mapOfStrings", Collections.singletonMap("key1", "val1"));
        updatedAgainSimpleEntity.setAttribute("arrayOfStrings", Arrays.asList("DummyTest3", "DummyTest4"));
        updatedAgainSimpleEntity.setAttribute("puArray", Collections.singletonList("1"));
        updatedAgainSimpleEntity.setAttribute("puMap", Collections.singletonMap("1", "1"));

        EntityMutationResponse updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), false);
        AtlasEntityHeader      updateHeader = updateRes.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
        AtlasEntity            updateEntity = getEntityFromStore(updateHeader);

        assertEquals(updateEntity.getAttribute("stringAtrr"), "Dummy String Test 3");
        assertEquals(updateEntity.getAttribute("arrayOfStrings"), Arrays.asList("DummyTest3", "DummyTest4"));
        assertEquals(updateEntity.getAttribute("mapOfStrings"), Collections.singletonMap("key1", "val1"));
        assertEquals(updateEntity.getAttribute("puArray"), Collections.singletonList("1"));
        assertEquals(updateEntity.getAttribute("puMap"), Collections.singletonMap("1", "1"));

        // full-update puArray and puMap; existing values should be replaced
        updatedAgainSimpleEntity.setAttribute("stringAtrr", "Dummy String Test 3");
        updatedAgainSimpleEntity.setAttribute("mapOfStrings", Collections.singletonMap("key1", "val1"));
        updatedAgainSimpleEntity.setAttribute("arrayOfStrings", Arrays.asList("DummyTest3", "DummyTest4"));
        updatedAgainSimpleEntity.setAttribute("puArray", Collections.singletonList("10"));
        updatedAgainSimpleEntity.setAttribute("puMap", Collections.singletonMap("10", "10"));

        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), false);
        updateHeader = updateRes.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
        updateEntity = getEntityFromStore(updateHeader);

        assertEquals(updateEntity.getAttribute("stringAtrr"), "Dummy String Test 3");
        assertEquals(updateEntity.getAttribute("arrayOfStrings"), Arrays.asList("DummyTest3", "DummyTest4"));
        assertEquals(updateEntity.getAttribute("mapOfStrings"), Collections.singletonMap("key1", "val1"));
        assertEquals(updateEntity.getAttribute("puArray"), Collections.singletonList("10"));
        assertEquals(updateEntity.getAttribute("puMap"), Collections.singletonMap("10", "10"));

        // partial-update tests
        // set puArray and puMap to null
        updatedAgainSimpleEntity.setAttribute("stringAtrr", "Dummy String Test 3");
        updatedAgainSimpleEntity.setAttribute("mapOfStrings", Collections.singletonMap("key1", "val1"));
        updatedAgainSimpleEntity.setAttribute("arrayOfStrings", Arrays.asList("DummyTest3", "DummyTest4"));
        updatedAgainSimpleEntity.setAttribute("puArray", null);
        updatedAgainSimpleEntity.setAttribute("puMap", null);

        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), false);
        updateHeader = updateRes.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
        updateEntity = getEntityFromStore(updateHeader);

        assertEquals(updateEntity.getAttribute("stringAtrr"), "Dummy String Test 3");
        assertEquals(updateEntity.getAttribute("arrayOfStrings"), Arrays.asList("DummyTest3", "DummyTest4"));
        assertEquals(updateEntity.getAttribute("mapOfStrings"), Collections.singletonMap("key1", "val1"));
        assertNull(updateEntity.getAttribute("puArray"));
        assertNull(updateEntity.getAttribute("puMap"));

        // partial-update: current value as null
        updatedAgainSimpleEntity.getAttributes().clear();
        updatedAgainSimpleEntity.setAttribute("puArray", Collections.singletonList("1"));
        updatedAgainSimpleEntity.setAttribute("puMap", Collections.singletonMap("1", "1"));

        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), true);
        updateHeader = updateRes.getFirstPartialUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
        updateEntity = getEntityFromStore(updateHeader);

        List<String>        puArray = new ArrayList<>(Collections.singletonList("1"));
        Map<String, String> puMap   = new HashMap<>(Collections.singletonMap("1", "1"));

        Assert.equals(updateEntity.getAttribute("puArray"), puArray);
        Assert.equals(updateEntity.getAttribute("puMap"), puMap);

        // partial-update: append to existing value
        updatedAgainSimpleEntity.getAttributes().clear();
        updatedAgainSimpleEntity.setAttribute("puArray", Collections.singletonList("2"));
        updatedAgainSimpleEntity.setAttribute("puMap", Collections.singletonMap("2", "2"));

        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), true);
        updateHeader = updateRes.getFirstPartialUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
        updateEntity = getEntityFromStore(updateHeader);

        puArray.add("2");
        puMap.put("2", "2");

        Assert.equals(updateEntity.getAttribute("puArray"), puArray);
        Assert.equals(updateEntity.getAttribute("puMap"), puMap);

        //  partial-update: with null value; existing value should be retained
        updatedAgainSimpleEntity.getAttributes().clear();
        updatedAgainSimpleEntity.setAttribute("puArray", null);
        updatedAgainSimpleEntity.setAttribute("puMap", null);
        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), true);
        updateHeader = updateRes.getFirstPartialUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
        updateEntity = getEntityFromStore(updateHeader);

        Assert.equals(updateEntity.getAttribute("puArray"), puArray);
        Assert.equals(updateEntity.getAttribute("puMap"), puMap);

        //  partial-update: with empty value; existing value should be retained
        updatedAgainSimpleEntity.getAttributes().clear();
        updatedAgainSimpleEntity.setAttribute("puArray", Collections.emptyList());
        updatedAgainSimpleEntity.setAttribute("puMap", Collections.emptyMap());

        updateRes    = entityStore.createOrUpdate(new AtlasEntityStream(updatedAgainSimpleEntity), true);
        updateHeader = updateRes.getFirstPartialUpdatedEntityByTypeName(ENTITY_TYPE_WITH_SIMPLE_ATTR);
        updateEntity = getEntityFromStore(updateHeader);

        Assert.equals(updateEntity.getAttribute("puArray"), puArray);
        Assert.equals(updateEntity.getAttribute("puMap"), puMap);
    }

    @Test(dependsOnMethods = "testCreateComplexAttributeEntity")
    public void testStructArray() throws Exception {
        init();
        AtlasEntity              complexEntity       = getEntityFromStore(complexCollectionAttrEntity.getEntity().getGuid());
        AtlasEntitiesWithExtInfo complexEntitiesInfo = new AtlasEntitiesWithExtInfo(complexEntity);

        // Modify array of structs
        List<AtlasStruct> structList = new ArrayList<>(Arrays.asList(new AtlasStruct("struct_type", "name", "structArray00"),
                new AtlasStruct("struct_type", "name", "structArray11"),
                new AtlasStruct("struct_type", "name", "structArray22")));

        complexEntity.setAttribute("listOfStructs", structList);

        EntityMutationResponse response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        AtlasEntityHeader      updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // add a new element to array of struct
        init();
        structList.add(new AtlasStruct("struct_type", "name", "structArray33"));
        complexEntity.setAttribute("listOfStructs", structList);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // remove one of the struct values - structArray00
        init();
        structList.remove(0);
        complexEntity.setAttribute("listOfStructs", structList);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // Update struct value within array of struct
        init();
        structList.get(0).setAttribute(NAME, "structArray11-edit");
        complexEntity.setAttribute("listOfStructs", structList);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // add a repeated element to array of struct
        init();
        structList.add(new AtlasStruct("struct_type", "name", "structArray33"));
        complexEntity.setAttribute("listOfStructs", structList);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // Remove all elements. Should set array attribute to null
        init();
        structList.clear();
        complexEntity.setAttribute("listOfStructs", structList);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));
    }

    @Test(dependsOnMethods = "testStructArray")
    public void testEntityArray() throws Exception {
        init();

        AtlasEntity              complexEntity       = getEntityFromStore(complexCollectionAttrEntity.getEntity().getGuid());
        AtlasEntitiesWithExtInfo complexEntitiesInfo = new AtlasEntitiesWithExtInfo(complexEntity);
        AtlasEntityType          entityType          = typeRegistry.getEntityTypeByName(ENTITY_TYPE);

        // Replace list of entities with new values
        AtlasEntity e0Array = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityArray00", "isReplicated", true));
        AtlasEntity e1Array = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityArray11", "isReplicated", false));
        AtlasEntity e2Array = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityArray22", "isReplicated", true));

        List<AtlasObjectId> entityList = new ArrayList<>(Arrays.asList(getAtlasObjectId(e0Array), getAtlasObjectId(e1Array), getAtlasObjectId(e2Array)));

        complexEntity.setAttribute("listOfEntities", entityList);
        complexEntitiesInfo.addReferredEntity(e0Array);
        complexEntitiesInfo.addReferredEntity(e1Array);
        complexEntitiesInfo.addReferredEntity(e2Array);

        init();
        EntityMutationResponse response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        AtlasEntityHeader      updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // add a new element to list of entities
        init();

        e0Array = entityStore.getByUniqueAttributes(entityType, getMap(NAME, "entityArray00", "isReplicated", true)).getEntity();
        e1Array = entityStore.getByUniqueAttributes(entityType, getMap(NAME, "entityArray11", "isReplicated", false)).getEntity();
        e2Array = entityStore.getByUniqueAttributes(entityType, getMap(NAME, "entityArray22", "isReplicated", true)).getEntity();

        AtlasEntity e3Array = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityArray33", "isReplicated", true));

        entityList = new ArrayList<>(Arrays.asList(getAtlasObjectId(e0Array), getAtlasObjectId(e1Array), getAtlasObjectId(e2Array), getAtlasObjectId(e3Array)));

        complexEntity.setAttribute("listOfEntities", entityList);
        complexEntitiesInfo.getReferredEntities().clear();
        complexEntitiesInfo.addReferredEntity(e3Array);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // remove one of the entity values - entityArray00
        init();
        e3Array    = entityStore.getByUniqueAttributes(entityType, getMap(NAME, "entityArray33", "isReplicated", true)).getEntity();
        entityList = new ArrayList<>(Arrays.asList(getAtlasObjectId(e1Array), getAtlasObjectId(e2Array), getAtlasObjectId(e3Array)));

        complexEntity.setAttribute("listOfEntities", entityList);
        complexEntitiesInfo.getReferredEntities().clear();

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // Update entity value within array of struct
        init();
        e1Array.setAttribute(NAME, "entityArray11-edit");
        complexEntity.setAttribute("listOfEntities", entityList);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // add a repeated element to array of struct
        init();
        AtlasEntity e3ArrayDuplicate = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityArray33", "isReplicated", true));

        entityList.add(getAtlasObjectId(e3ArrayDuplicate));
        complexEntity.setAttribute("listOfEntities", entityList);
        complexEntitiesInfo.getReferredEntities().clear();
        complexEntitiesInfo.addReferredEntity(e3ArrayDuplicate);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // Remove all elements. Should set array attribute to null
        init();
        entityList.clear();
        complexEntity.setAttribute("listOfEntities", entityList);
        complexEntitiesInfo.getReferredEntities().clear();

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));
    }

    @Test(dependsOnMethods = "testEntityArray")
    public void testStructMap() throws Exception {
        init();
        AtlasEntity              complexEntity       = getEntityFromStore(complexCollectionAttrEntity.getEntity().getGuid());
        AtlasEntitiesWithExtInfo complexEntitiesInfo = new AtlasEntitiesWithExtInfo(complexEntity);

        // Modify map of structs
        Map<String, AtlasStruct> structMap = getMap("key00", new AtlasStruct("struct_type", "name", "structMap00"),
                "key11", new AtlasStruct("struct_type", "name", "structMap11"),
                "key22", new AtlasStruct("struct_type", "name", "structMap22"));

        complexEntity.setAttribute("mapOfStructs", structMap);

        EntityMutationResponse response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        AtlasEntityHeader      updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // add a new element to map of struct - structMap6
        init();
        structMap.put("key33", new AtlasStruct("struct_type", "name", "structMap33"));
        complexEntity.setAttribute("mapOfStructs", structMap);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // remove one of the entity values - structMap
        init();
        structMap.remove("key00");
        complexEntity.setAttribute("mapOfStructs", structMap);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // Update struct value within map of struct
        init();
        structMap.get("key11").setAttribute("name", "structMap11-edit");
        complexEntity.setAttribute("mapOfStructs", structMap);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // add a repeated element to array of struct
        init();
        structMap.put("key33", new AtlasStruct("struct_type", "name", "structMap33"));
        complexEntity.setAttribute("mapOfStructs", structMap);
        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        // no update since duplicate entry
        assertNull(updatedComplexEntity);

        // Remove all elements. Should set array attribute to null
        init();
        structMap.clear();
        complexEntity.setAttribute("mapOfStructs", structMap);
        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));
    }

    @Test(dependsOnMethods = "testStructMap")
    public void testEntityMap() throws Exception {
        init();
        AtlasEntity              complexEntity       = getEntityFromStore(complexCollectionAttrEntity.getEntity().getGuid());
        AtlasEntitiesWithExtInfo complexEntitiesInfo = new AtlasEntitiesWithExtInfo(complexEntity);

        // Modify map of entities
        AtlasEntity e0MapValue = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityMapValue00", "isReplicated", false));
        AtlasEntity e1MapValue = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityMapValue11", "isReplicated", true));
        AtlasEntity e2MapValue = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityMapValue22", "isReplicated", false));
        Map<String, Object> entityMap = getMap("key00", getAtlasObjectId(e0MapValue), "key11", getAtlasObjectId(e1MapValue), "key22", getAtlasObjectId(e2MapValue));

        complexEntity.setAttribute("mapOfEntities", entityMap);
        complexEntitiesInfo.addReferredEntity(e0MapValue);
        complexEntitiesInfo.addReferredEntity(e1MapValue);
        complexEntitiesInfo.addReferredEntity(e2MapValue);

        init();

        EntityMutationResponse response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        AtlasEntityHeader      updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // add a new element to map of entities
        init();

        AtlasEntity e3MapValue = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityMapValue33", "isReplicated", false));

        entityMap.put("key33", getAtlasObjectId(e3MapValue));
        complexEntity.setAttribute("mapOfEntities", entityMap);
        complexEntitiesInfo.addReferredEntity(e3MapValue);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // remove one of the entity values - [key00 : entityMapValue00]
        init();
        entityMap.remove("key00");
        complexEntity.setAttribute("mapOfEntities", entityMap);
        complexEntitiesInfo.addReferredEntity(e3MapValue);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // Update entity value within map of entities
        init();

        AtlasEntity e1MapValueEdit = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityMapValue11-edit", "isReplicated", false));

        entityMap.clear();
        entityMap.put("key11", getAtlasObjectId(e1MapValueEdit));
        entityMap.put("key22", getAtlasObjectId(e2MapValue));
        entityMap.put("key33", getAtlasObjectId(e3MapValue));
        complexEntity.setAttribute("mapOfEntities", entityMap);
        complexEntitiesInfo.addReferredEntity(e1MapValueEdit);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // add a repeated element to map of entities
        init();

        e3MapValue = new AtlasEntity(ENTITY_TYPE, getMap(NAME, "entityMapValue33", "isReplicated", false));

        entityMap.put("key33", getAtlasObjectId(e3MapValue));
        complexEntity.setAttribute("mapOfEntities", entityMap);
        complexEntitiesInfo.addReferredEntity(e3MapValue);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));

        // Remove all elements. Should set map attribute to null
        init();
        entityMap.clear();
        complexEntity.setAttribute("mapOfEntities", entityMap);
        complexEntitiesInfo.addReferredEntity(e3MapValue);

        response             = entityStore.createOrUpdate(new AtlasEntityStream(complexEntitiesInfo), false);
        updatedComplexEntity = response.getFirstUpdatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);
        validateEntity(complexEntitiesInfo, getEntityFromStore(updatedComplexEntity));
    }

    @Test(dependsOnMethods = "testEntityMap")
    public void testDeleteEntityRemoveReferences() throws Exception {
        init();

        complexCollectionAttrEntityForDelete.getEntity().setAttribute(NAME, ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR_DELETE);

        EntityMutationResponse response      = entityStore.createOrUpdate(new AtlasEntityStream(complexCollectionAttrEntityForDelete), false);
        AtlasEntityHeader      entityCreated = response.getFirstCreatedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        validateEntity(complexCollectionAttrEntityForDelete, getEntityFromStore(entityCreated));

        // delete entity and check if referenced complex attribute edges are deleted
        response = entityStore.deleteById(entityCreated.getGuid());

        AtlasEntityHeader entityDeleted = response.getFirstDeletedEntityByTypeName(ENTITY_TYPE_WITH_COMPLEX_COLLECTION_ATTR);

        GraphTransactionInterceptor.clearCache();
        AtlasEntityWithExtInfo deletedEntityWithExtInfo = entityStore.getById(entityDeleted.getGuid());
        AtlasVertex            deletedEntityVertex      = AtlasGraphUtilsV2.findByGuid(entityDeleted.getGuid());
        Iterator<AtlasEdge>    edges                    = deletedEntityVertex.getEdges(AtlasEdgeDirection.OUT).iterator();

        // validate all attribute edges are deleted
        while (edges != null && edges.hasNext()) {
            assertEquals(getStatus(edges.next()), AtlasEntity.Status.DELETED);
        }

        AtlasEntity                deletedEntity  = deletedEntityWithExtInfo.getEntity();
        List<AtlasObjectId>        listOfEntities = (List<AtlasObjectId>) deletedEntity.getAttribute("listOfEntities");
        Map<String, AtlasObjectId> mapOfEntities  = (Map<String, AtlasObjectId>) deletedEntity.getAttribute("mapOfEntities");

        // validate entity attributes are deleted
        for (AtlasObjectId o : listOfEntities) {
            AtlasEntity entity = deletedEntityWithExtInfo.getEntity(o.getGuid());
            assertEquals(entity.getStatus(), AtlasEntity.Status.DELETED);
        }

        for (AtlasObjectId o : mapOfEntities.values()) {
            AtlasEntity entity = deletedEntityWithExtInfo.getEntity(o.getGuid());
            assertEquals(entity.getStatus(), AtlasEntity.Status.DELETED);
        }
    }

    private void updateEntityMapAttributes(AtlasEntity attrEntity, Map<String, String> map1, Map<String, Integer> map2, Map<String, Boolean> map3, Map<String, Float> map4, Map<String, Date> map5) {
        attrEntity.setAttribute("mapAttr1", map1);
        attrEntity.setAttribute("mapAttr2", map2);
        attrEntity.setAttribute("mapAttr3", map3);
        attrEntity.setAttribute("mapAttr4", map4);
        attrEntity.setAttribute("mapAttr5", map5);
    }

    private static <K, V> Map<K, V> getMap(K key1, V val1, K key2, V val2) {
        Map<K, V> ret = new HashMap<>();

        ret.put(key1, val1);
        ret.put(key2, val2);

        return ret;
    }

    private static <K, V> Map<K, V> getMap(K key1, V val1, K key2, V val2, K key3, V val3) {
        Map<K, V> ret = new HashMap<>();

        ret.put(key1, val1);
        ret.put(key2, val2);
        ret.put(key3, val3);

        return ret;
    }
}
