Java Code Examples for org.apache.atlas.model.instance.EntityMutationResponse#getFirstUpdatedEntityByTypeName()

@Test(dependsOnMethods = "testCreate")
public void testMapOfPrimitivesUpdate() throws Exception {
    AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
    AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);

    entitiesInfo.addReferredEntity(tableEntity);

    //Add a new entry
    Map<String, String> paramsMap = (Map<String, String>) tableEntity.getAttribute("parametersMap");
    paramsMap.put("newParam", "value");
    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 1);
    AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //Remove an entry
    paramsMap.remove("key1");
    tableEntity.setAttribute("parametersMap", paramsMap);
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 1);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
}
 

@Test(dependsOnMethods = "testCreate")
public void testClassUpdate() throws Exception {

    init();
    //Create new db instance
    final AtlasEntity databaseInstance = TestUtilsV2.createDBEntity();

    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(databaseInstance), false);
    final AtlasEntityHeader dbCreated = response.getFirstCreatedEntityByTypeName(TestUtilsV2.DATABASE_TYPE);

    init();
    Map<String, AtlasEntity> tableCloneMap = new HashMap<>();
    AtlasEntity tableClone = new AtlasEntity(tblEntity.getEntity());
    tableClone.setAttribute("database", new AtlasObjectId(dbCreated.getGuid(), TestUtilsV2.DATABASE_TYPE));

    tableCloneMap.put(dbCreated.getGuid(), databaseInstance);
    tableCloneMap.put(tableClone.getGuid(), tableClone);

    response = entityStore.createOrUpdate(new InMemoryMapEntityStream(tableCloneMap), false);
    final AtlasEntityHeader tableDefinition = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    AtlasEntity updatedTableDefinition = getEntityFromStore(tableDefinition);
    assertNotNull(updatedTableDefinition.getAttribute("database"));
    Assert.assertEquals(((AtlasObjectId) updatedTableDefinition.getAttribute("database")).getGuid(), dbCreated.getGuid());
}
 

@Test(dependsOnMethods = "testCreate")
public void testMapOfPrimitivesUpdate() throws Exception {
    AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
    AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);

    entitiesInfo.addReferredEntity(tableEntity);

    //Add a new entry
    Map<String, String> paramsMap = (Map<String, String>) tableEntity.getAttribute("parametersMap");
    paramsMap.put("newParam", "value");
    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 1);
    AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //Remove an entry
    paramsMap.remove("key1");
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    validateMutationResponse(response, EntityMutations.EntityOperation.UPDATE, 1);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
}
 

@Test(dependsOnMethods = "testCreate")
public void testClassUpdate() throws Exception {

    init();
    //Create new db instance
    final AtlasEntity databaseInstance = TestUtilsV2.createDBEntity();

    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(databaseInstance), false);
    final AtlasEntityHeader dbCreated = response.getFirstCreatedEntityByTypeName(TestUtilsV2.DATABASE_TYPE);

    init();
    Map<String, AtlasEntity> tableCloneMap = new HashMap<>();
    AtlasEntity tableClone = new AtlasEntity(tblEntity.getEntity());
    tableClone.setAttribute("database", new AtlasObjectId(dbCreated.getGuid(), TestUtils.DATABASE_TYPE));

    tableCloneMap.put(dbCreated.getGuid(), databaseInstance);
    tableCloneMap.put(tableClone.getGuid(), tableClone);

    response = entityStore.createOrUpdate(new InMemoryMapEntityStream(tableCloneMap), false);
    final AtlasEntityHeader tableDefinition = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    AtlasEntity updatedTableDefinition = getEntityFromStore(tableDefinition);
    Assert.assertNotNull(updatedTableDefinition.getAttribute("database"));
    Assert.assertEquals(((AtlasObjectId) updatedTableDefinition.getAttribute("database")).getGuid(), dbCreated.getGuid());
}
 

@Test(dependsOnMethods = "testCreate")
public void testStructs() throws Exception {
    AtlasEntity              databaseEntity = dbEntity.getEntity();
    AtlasEntity              tableEntity    = new AtlasEntity(tblEntity.getEntity());
    AtlasEntitiesWithExtInfo entitiesInfo   = new AtlasEntitiesWithExtInfo(tableEntity);

    AtlasStruct serdeInstance = new AtlasStruct(TestUtilsV2.SERDE_TYPE, NAME, "serde1Name");
    serdeInstance.setAttribute("serde", "test");
    serdeInstance.setAttribute("description", "testDesc");
    tableEntity.setAttribute("serde1", serdeInstance);
    tableEntity.setAttribute("database", new AtlasObjectId(databaseEntity.getTypeName(), NAME, databaseEntity.getAttribute(NAME)));

    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //update struct attribute
    serdeInstance.setAttribute("serde", "testUpdated");
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //set to null
    tableEntity.setAttribute("description", null);
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    Assert.assertNull(updatedTable.getAttribute("description"));
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
}
 

@Test(dependsOnMethods = "testCreate")
public void testStructs() throws Exception {
    AtlasEntity              databaseEntity = dbEntity.getEntity();
    AtlasEntity              tableEntity    = new AtlasEntity(tblEntity.getEntity());
    AtlasEntitiesWithExtInfo entitiesInfo   = new AtlasEntitiesWithExtInfo(tableEntity);

    AtlasStruct serdeInstance = new AtlasStruct(TestUtils.SERDE_TYPE, TestUtilsV2.NAME, "serde1Name");
    serdeInstance.setAttribute("serde", "test");
    serdeInstance.setAttribute("description", "testDesc");
    tableEntity.setAttribute("serde1", serdeInstance);
    tableEntity.setAttribute("database", new AtlasObjectId(databaseEntity.getTypeName(), TestUtilsV2.NAME, databaseEntity.getAttribute(TestUtilsV2.NAME)));

    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //update struct attribute
    serdeInstance.setAttribute("serde", "testUpdated");
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //set to null
    tableEntity.setAttribute("description", null);
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    Assert.assertNull(updatedTable.getAttribute("description"));
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
}
 

@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<String, String>() {{ put("map1Key11", "value11");
                                                                put("map1Key22", "value22");
                                                                put("map1Key33", "value33"); }};

    Map<String, Integer> map2 = new HashMap<String, Integer>() {{ put("map2Key11", 1100);
                                                                  put("map2Key22", 2200);
                                                                  put("map2Key33", 3300); }};

    Map<String, Boolean> map3 = new HashMap<String, Boolean>() {{ put("map3Key11", true);
                                                                  put("map3Key22", false);
                                                                  put("map3Key33", true); }};

    Map<String, Float> map4 = new HashMap<String, Float>() {{ put("map4Key11", 11.0f);
                                                              put("map4Key22", 22.0f);
                                                              put("map4Key33", 33.0f); }};

    Map<String, Date> map5 = new HashMap<String, Date>() {{ put("map5Key11", DateUtils.addHours(new Date(), 1));
                                                            put("map5Key22", DateUtils.addHours(new Date(), 2));
                                                            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(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, new HashMap<String, Object>() {{ put(NAME, "entityArray00"); put("isReplicated", true); }});
    AtlasEntity e1Array = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityArray11"); put("isReplicated", false); }});
    AtlasEntity e2Array = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityArray22"); put("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, new HashMap<String, Object>() {{ put(NAME, "entityArray00"); put("isReplicated", true); }}).getEntity();
    e1Array = entityStore.getByUniqueAttributes(entityType, new HashMap<String, Object>() {{ put(NAME, "entityArray11"); put("isReplicated", false); }}).getEntity();
    e2Array = entityStore.getByUniqueAttributes(entityType, new HashMap<String, Object>() {{ put(NAME, "entityArray22"); put("isReplicated", true); }}).getEntity();
    AtlasEntity e3Array = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityArray33"); put("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, new HashMap<String, Object>() {{ put(NAME, "entityArray33"); put("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 e3Array_duplicate = new AtlasEntity(ENTITY_TYPE, new HashMap<String, Object>() {{ put(NAME, "entityArray33"); put("isReplicated", true); }});
    entityList.add(getAtlasObjectId(e3Array_duplicate));
    complexEntity.setAttribute("listOfEntities", entityList);
    complexEntitiesInfo.getReferredEntities().clear();
    complexEntitiesInfo.addReferredEntity(e3Array_duplicate);

    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
    HashMap<String, AtlasStruct> structMap = new HashMap<String, AtlasStruct>() {{
                                                    put("key00", new AtlasStruct("struct_type", "name", "structMap00"));
                                                    put("key11", new AtlasStruct("struct_type", "name", "structMap11"));
                                                    put("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 = "testCreate")
public void testArrayOfEntityUpdate() throws Exception {
    AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
    List<AtlasObjectId>      columns      = new ArrayList<>();
    AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);

    AtlasEntity col1 = TestUtilsV2.createColumnEntity(tableEntity);
    col1.setAttribute(NAME, "col1");

    AtlasEntity col2 = TestUtilsV2.createColumnEntity(tableEntity);
    col2.setAttribute(NAME, "col2");

    columns.add(AtlasTypeUtil.getAtlasObjectId(col1));
    columns.add(AtlasTypeUtil.getAtlasObjectId(col2));

    tableEntity.setAttribute(COLUMNS_ATTR_NAME, columns);

    entitiesInfo.addReferredEntity(dbEntity.getEntity());
    entitiesInfo.addReferredEntity(col1);
    entitiesInfo.addReferredEntity(col2);

    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);

    AtlasEntityHeader updatedTableHeader = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
    validateEntity(entitiesInfo, getEntityFromStore(updatedTableHeader));

    //Complete update. Add  array elements - col3,col4
    AtlasEntity col3 = TestUtilsV2.createColumnEntity(tableEntity);
    col3.setAttribute(NAME, "col3");

    AtlasEntity col4 = TestUtilsV2.createColumnEntity(tableEntity);
    col4.setAttribute(NAME, "col4");

    columns.add(AtlasTypeUtil.getAtlasObjectId(col3));
    columns.add(AtlasTypeUtil.getAtlasObjectId(col4));
    tableEntity.setAttribute(COLUMNS_ATTR_NAME, columns);

    entitiesInfo.addReferredEntity(col3);
    entitiesInfo.addReferredEntity(col4);

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

    updatedTableHeader = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
    validateEntity(entitiesInfo, getEntityFromStore(updatedTableHeader));

    //Swap elements
    columns.clear();
    columns.add(AtlasTypeUtil.getAtlasObjectId(col4));
    columns.add(AtlasTypeUtil.getAtlasObjectId(col3));
    tableEntity.setAttribute(COLUMNS_ATTR_NAME, columns);

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

    updatedTableHeader = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
    AtlasEntity updatedEntity = getEntityFromStore(updatedTableHeader);
    // deleted columns are also included in "columns" attribute
    Assert.assertTrue(((List<AtlasObjectId>) updatedEntity.getAttribute(COLUMNS_ATTR_NAME)).size() >= 2);

    assertEquals(response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 2);  // col1, col2 are deleted

    //Update array column to null
    tableEntity.setAttribute(COLUMNS_ATTR_NAME, null);

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

    updatedTableHeader = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
    validateEntity(entitiesInfo, getEntityFromStore(updatedTableHeader));
    assertEquals(response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 2);
}
 

@Test(dependsOnMethods = "testCreate")
public void testArrayOfStructs() throws Exception {
    //Modify array of structs
    AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
    AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);

    List<AtlasStruct> partitions = new ArrayList<AtlasStruct>(){{
        add(new AtlasStruct(TestUtilsV2.PARTITION_STRUCT_TYPE, NAME, "part1"));
        add(new AtlasStruct(TestUtilsV2.PARTITION_STRUCT_TYPE, NAME, "part2"));
    }};
    tableEntity.setAttribute("partitions", partitions);

    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //add a new element to array of struct
    partitions.add(new AtlasStruct(TestUtilsV2.PARTITION_STRUCT_TYPE, NAME, "part3"));
    tableEntity.setAttribute("partitions", partitions);
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //remove one of the struct values
    partitions.remove(1);
    tableEntity.setAttribute("partitions", partitions);
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //Update struct value within array of struct
    partitions.get(0).setAttribute(NAME, "part4");
    tableEntity.setAttribute("partitions", partitions);
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));


    //add a repeated element to array of struct
    partitions.add(new AtlasStruct(TestUtilsV2.PARTITION_STRUCT_TYPE, NAME, "part4"));
    tableEntity.setAttribute("partitions", partitions);
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    // Remove all elements. Should set array attribute to null
    partitions.clear();
    tableEntity.setAttribute("partitions", partitions);
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
}
 

@Test(dependsOnMethods = "testCreate")
public void testArrayOfEntityUpdate() throws Exception {
    AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
    List<AtlasObjectId>      columns      = new ArrayList<>();
    AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);


    AtlasEntity col1 = TestUtilsV2.createColumnEntity(tableEntity);
    col1.setAttribute(TestUtilsV2.NAME, "col1");

    AtlasEntity col2 = TestUtilsV2.createColumnEntity(tableEntity);
    col2.setAttribute(TestUtilsV2.NAME, "col2");

    columns.add(AtlasTypeUtil.getAtlasObjectId(col1));
    columns.add(AtlasTypeUtil.getAtlasObjectId(col2));

    tableEntity.setAttribute(TestUtilsV2.COLUMNS_ATTR_NAME, columns);

    entitiesInfo.addReferredEntity(dbEntity.getEntity());
    entitiesInfo.addReferredEntity(col1);
    entitiesInfo.addReferredEntity(col2);

    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);

    AtlasEntityHeader updatedTableHeader = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
    validateEntity(entitiesInfo, getEntityFromStore(updatedTableHeader));

    //Complete update. Add  array elements - col3,col4
    AtlasEntity col3 = TestUtilsV2.createColumnEntity(tableEntity);
    col3.setAttribute(TestUtilsV2.NAME, "col3");

    AtlasEntity col4 = TestUtilsV2.createColumnEntity(tableEntity);
    col4.setAttribute(TestUtilsV2.NAME, "col4");

    columns.add(AtlasTypeUtil.getAtlasObjectId(col3));
    columns.add(AtlasTypeUtil.getAtlasObjectId(col4));
    tableEntity.setAttribute(TestUtilsV2.COLUMNS_ATTR_NAME, columns);

    entitiesInfo.addReferredEntity(col3);
    entitiesInfo.addReferredEntity(col4);

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

    updatedTableHeader = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
    validateEntity(entitiesInfo, getEntityFromStore(updatedTableHeader));

    //Swap elements
    columns.clear();
    columns.add(AtlasTypeUtil.getAtlasObjectId(col4));
    columns.add(AtlasTypeUtil.getAtlasObjectId(col3));
    tableEntity.setAttribute(TestUtilsV2.COLUMNS_ATTR_NAME, columns);

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

    updatedTableHeader = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
    AtlasEntity updatedEntity = getEntityFromStore(updatedTableHeader);
    // deleted columns are also included in "columns" attribute
    Assert.assertTrue(((List<AtlasObjectId>) updatedEntity.getAttribute(COLUMNS_ATTR_NAME)).size() >= 2);

    assertEquals(response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 2);  // col1, col2 are deleted

    //Update array column to null
    tableEntity.setAttribute(COLUMNS_ATTR_NAME, null);

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

    updatedTableHeader = response.getFirstUpdatedEntityByTypeName(tableEntity.getTypeName());
    validateEntity(entitiesInfo, getEntityFromStore(updatedTableHeader));
    assertEquals(response.getEntitiesByOperation(EntityMutations.EntityOperation.DELETE).size(), 2);
}
 

@Test(dependsOnMethods = "testCreate")
public void testArrayOfStructs() throws Exception {
    //Modify array of structs
    AtlasEntity              tableEntity  = new AtlasEntity(tblEntity.getEntity());
    AtlasEntitiesWithExtInfo entitiesInfo = new AtlasEntitiesWithExtInfo(tableEntity);

    List<AtlasStruct> partitions = new ArrayList<AtlasStruct>(){{
        add(new AtlasStruct(TestUtilsV2.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "part1"));
        add(new AtlasStruct(TestUtilsV2.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "part2"));
    }};
    tableEntity.setAttribute("partitions", partitions);

    init();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    AtlasEntityHeader updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //add a new element to array of struct
    partitions.add(new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "part3"));
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //remove one of the struct values
    init();
    partitions.remove(1);
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    //Update struct value within array of struct
    init();
    partitions.get(0).setAttribute(TestUtilsV2.NAME, "part4");
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));


    //add a repeated element to array of struct
    partitions.add(new AtlasStruct(TestUtils.PARTITION_STRUCT_TYPE, TestUtilsV2.NAME, "part4"));
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));

    // Remove all elements. Should set array attribute to null
    partitions.clear();
    init();
    response = entityStore.createOrUpdate(new AtlasEntityStream(entitiesInfo), false);
    updatedTable = response.getFirstUpdatedEntityByTypeName(TABLE_TYPE);
    validateEntity(entitiesInfo, getEntityFromStore(updatedTable));
}
 

@Test
public void testUpdateEntity_MultiplicityOneNonCompositeReference() throws Exception {
    AtlasEntity.AtlasEntitiesWithExtInfo hrDept = TestUtilsV2.createDeptEg2();
    init();

    RequestContextV1.clear();
    final EntityMutationResponse hrDeptCreationResponse = entityStore.createOrUpdate(new AtlasEntityStream(hrDept), false);
    final AtlasEntityHeader deptCreated = hrDeptCreationResponse.getFirstUpdatedEntityByTypeName(DEPARTMENT_TYPE);
    final AtlasEntityHeader maxEmployeeCreated = hrDeptCreationResponse.getCreatedEntityByTypeNameAndAttribute(TestUtilsV2.EMPLOYEE_TYPE, NAME, "Max");
    final AtlasEntityHeader johnEmployeeCreated = hrDeptCreationResponse.getUpdatedEntityByTypeNameAndAttribute(TestUtilsV2.EMPLOYEE_TYPE, NAME, "John");
    final AtlasEntityHeader janeEmployeeCreated = hrDeptCreationResponse.getCreatedEntityByTypeNameAndAttribute(TestUtilsV2.MANAGER_TYPE, NAME, "Jane");
    final AtlasEntityHeader juliusEmployeeCreated = hrDeptCreationResponse.getUpdatedEntityByTypeNameAndAttribute(TestUtilsV2.MANAGER_TYPE, NAME, "Julius");

    ITypedReferenceableInstance max = metadataService.getEntityDefinition(maxEmployeeCreated.getGuid());
    String maxGuid = max.getId()._getId();
    AtlasVertex vertex = GraphHelper.getInstance().getVertexForGUID(maxGuid);
    Long creationTimestamp = GraphHelper.getSingleValuedProperty(vertex, Constants.TIMESTAMP_PROPERTY_KEY, Long.class);
    Assert.assertNotNull(creationTimestamp);

    Long modificationTimestampPreUpdate = GraphHelper.getSingleValuedProperty(vertex, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY, Long.class);
    Assert.assertNotNull(modificationTimestampPreUpdate);

    AtlasEntity maxEmployee = getEmployeeByName(hrDept, "Max");
    maxEmployee.setAttribute("mentor", AtlasTypeUtil.getAtlasObjectId(johnEmployeeCreated));
    maxEmployee.setAttribute("department", AtlasTypeUtil.getAtlasObjectId(deptCreated));
    maxEmployee.setAttribute("manager", AtlasTypeUtil.getAtlasObjectId(janeEmployeeCreated));

    init();
    EntityMutationResponse entityResult = entityStore.createOrUpdate(new AtlasEntityStream(maxEmployee), false);

    assertEquals(entityResult.getUpdatedEntities().size(), 1);
    assertTrue(extractGuids(entityResult.getUpdatedEntities()).contains(maxGuid));

    // Verify the update was applied correctly - john should now be max's mentor.
    max = metadataService.getEntityDefinition(maxGuid);
    ITypedReferenceableInstance refTarget = (ITypedReferenceableInstance) max.get("mentor");
    Assert.assertEquals(refTarget.getId()._getId(), johnEmployeeCreated.getGuid());

    // Verify modification timestamp was updated.
    vertex = GraphHelper.getInstance().getVertexForGUID(maxGuid);
    Long modificationTimestampPostUpdate = GraphHelper.getSingleValuedProperty(vertex, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY, Long.class);
    Assert.assertNotNull(modificationTimestampPostUpdate);
    Assert.assertTrue(creationTimestamp < modificationTimestampPostUpdate);

    // Update max's mentor reference to jane.
    maxEmployee.setAttribute("mentor", AtlasTypeUtil.getAtlasObjectId(janeEmployeeCreated));
    init();
    entityResult = entityStore.createOrUpdate(new AtlasEntityStream(maxEmployee), false);
    assertEquals(entityResult.getUpdatedEntities().size(), 1);
    assertTrue(extractGuids(entityResult.getUpdatedEntities()).contains(maxGuid));

    // Verify the update was applied correctly - jane should now be max's mentor.
    max = metadataService.getEntityDefinition(maxGuid);
    refTarget = (ITypedReferenceableInstance) max.get("mentor");
    Assert.assertEquals(refTarget.getId()._getId(), janeEmployeeCreated.getGuid());

    // Verify modification timestamp was updated.
    vertex = GraphHelper.getInstance().getVertexForGUID(maxGuid);
    Long modificationTimestampPost2ndUpdate = GraphHelper.getSingleValuedProperty(vertex, Constants.MODIFICATION_TIMESTAMP_PROPERTY_KEY, Long.class);
    Assert.assertNotNull(modificationTimestampPost2ndUpdate);
    Assert.assertTrue(modificationTimestampPostUpdate < modificationTimestampPost2ndUpdate);

    ITypedReferenceableInstance julius = metadataService.getEntityDefinition(juliusEmployeeCreated.getGuid());
    Id juliusId = julius.getId();

    init();
    maxEmployee.setAttribute("manager", AtlasTypeUtil.getAtlasObjectId(juliusEmployeeCreated));
    entityResult = entityStore.createOrUpdate(new AtlasEntityStream(maxEmployee), false);
    assertEquals(entityResult.getUpdatedEntities().size(), 3);
    List<String> updatedGuids = extractGuids(entityResult.getUpdatedEntities());
    assertTrue(updatedGuids.contains(maxGuid));
    assertTrue(updatedGuids.contains(janeEmployeeCreated.getGuid()));
    // Should have updated julius to add max in subordinates list.
    assertTrue(updatedGuids.contains(juliusEmployeeCreated.getGuid()));

    // Verify the update was applied correctly - julius should now be max's manager and max should be julius' subordinate.
    max = metadataService.getEntityDefinition(maxGuid);
    refTarget = (ITypedReferenceableInstance) max.get("manager");
    Assert.assertEquals(refTarget.getId()._getId(), juliusId._getId());
    julius = metadataService.getEntityDefinition(juliusId._getId());
    Object value = julius.get("subordinates");
    Assert.assertTrue(value instanceof List);
    List<ITypedReferenceableInstance> refList = (List<ITypedReferenceableInstance>) value;
    Assert.assertEquals(refList.size(), 1);
    Assert.assertEquals(refList.get(0).getId()._getId(), maxGuid);

    assertTestUpdateEntity_MultiplicityOneNonCompositeReference(janeEmployeeCreated.getGuid());
}