org.apache.atlas.model.typedef.AtlasTypesDef Java Examples

public void addTypes(AtlasTypesDef typesDef) throws AtlasBaseException {
    if (LOG.isDebugEnabled()) {
        LOG.debug("==> AtlasTypeRegistry.addTypes({})", typesDef);
    }

    if (typesDef != null) {
        addTypesWithNoRefResolve(typesDef.getEnumDefs());
        addTypesWithNoRefResolve(typesDef.getStructDefs());
        addTypesWithNoRefResolve(typesDef.getClassificationDefs());
        addTypesWithNoRefResolve(typesDef.getEntityDefs());
        addTypesWithNoRefResolve(typesDef.getRelationshipDefs());

        resolveReferences();
    }

    if (LOG.isDebugEnabled()) {
        LOG.debug("<== AtlasTypeRegistry.addTypes({})", typesDef);
    }
}
 

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

    AtlasTypesDef[] testTypesDefs = new AtlasTypesDef[] { TestUtilsV2.defineDeptEmployeeTypes(),
                                                          TestUtilsV2.defineHiveTypes(),
                                                          TestUtilsV2.defineTypeWithNestedCollectionAttributes(),
                                                          TestUtilsV2.defineEnumTypes(),
                                                          TestUtilsV2.defineBusinessMetadataTypes()
                                                        };
    createTypesDef(testTypesDefs);

    deptEntity                 = TestUtilsV2.createDeptEg2();
    dbEntity                   = TestUtilsV2.createDBEntityV2();
    tblEntity                  = TestUtilsV2.createTableEntityV2(dbEntity.getEntity());
    nestedCollectionAttrEntity = TestUtilsV2.createNestedCollectionAttrEntity();
    primitiveEntity            = TestUtilsV2.createprimitiveEntityV2();

    AtlasTypesDef typesDef11         = new AtlasTypesDef();
    List          primitiveEntityDef = new ArrayList<AtlasEntityDef>();

    primitiveEntityDef.add(TestUtilsV2.createPrimitiveEntityDef());
    typesDef11.setEntityDefs(primitiveEntityDef);

    typeDefStore.createTypesDef(typesDef11);
}
 

public static AtlasTypesDef simpleType(){
    AtlasEntityDef superTypeDefinition =
            createClassTypeDef("h_type", Collections.<String>emptySet(),
                    createOptionalAttrDef("attr", "string"));

    AtlasStructDef structTypeDefinition = new AtlasStructDef("s_type", "structType", "1.0",
            Arrays.asList(createRequiredAttrDef("name", "string")));

    AtlasClassificationDef traitTypeDefinition =
            AtlasTypeUtil.createTraitTypeDef("t_type", "traitType", Collections.<String>emptySet());

    AtlasEnumDef enumTypeDefinition = new AtlasEnumDef("e_type", "enumType", "1.0",
            Arrays.asList(new AtlasEnumElementDef("ONE", "Element Description", 1)));

    AtlasTypesDef ret = AtlasTypeUtil.getTypesDef(Collections.singletonList(enumTypeDefinition), Collections.singletonList(structTypeDefinition),
            Collections.singletonList(traitTypeDefinition), Collections.singletonList(superTypeDefinition));

    populateSystemAttributes(ret);

    return ret;
}
 

@BeforeClass
public void setUp() throws Exception {
    new GraphBackedSearchIndexer(typeRegistry);

    // create employee relationship types
    AtlasTypesDef employeeTypes = getDepartmentEmployeeTypes();
    typeDefStore.createTypesDef(employeeTypes);

    AtlasEntitiesWithExtInfo employeeInstances = getDepartmentEmployeeInstances();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(employeeInstances), false);

    for (AtlasEntityHeader entityHeader : response.getCreatedEntities()) {
        employeeNameIdMap.put((String) entityHeader.getAttribute(NAME), getAtlasObjectId(entityHeader));
    }

    init();
    AtlasTypesDef testTypes = getInverseReferenceTestTypes();
    typeDefStore.createTypesDef(testTypes);
}
 

@Test(dataProvider = "zipFileStocks")
public void examineContents_BehavesAsExpected(ZipSource zipSource) throws IOException, AtlasBaseException {
    List<String> creationOrder = zipSource.getCreationOrder();

    assertNotNull(creationOrder);
    assertEquals(creationOrder.size(), 4);

    AtlasTypesDef typesDef = zipSource.getTypesDef();
    assertNotNull(typesDef);
    assertEquals(typesDef.getEntityDefs().size(), 6);

    useCreationOrderToFetchEntitiesWithExtInfo(zipSource, creationOrder);
    useCreationOrderToFetchEntities(zipSource, creationOrder);
    attemptToFetchNonExistentGuid_ReturnsNull(zipSource, "non-existent-guid");
    verifyGuidRemovalOnImportComplete(zipSource, creationOrder.get(0));
}
 

@Test
public void testDuplicateCreate() throws Exception {
    AtlasEntityDef type = createClassTypeDef(randomString(),
            Collections.<String>emptySet(), AtlasTypeUtil.createUniqueRequiredAttrDef("name", "string"));
    AtlasTypesDef typesDef = new AtlasTypesDef();
    typesDef.getEntityDefs().add(type);

    AtlasTypesDef created = clientV2.createAtlasTypeDefs(typesDef);
    assertNotNull(created);

    try {
        created = clientV2.createAtlasTypeDefs(typesDef);
        fail("Expected 409");
    } catch (AtlasServiceException e) {
        assertEquals(e.getStatus().getStatusCode(), Response.Status.CONFLICT.getStatusCode());
    }
}
 

@Test
public void testUTF8() throws Exception {
    String classType = randomString();
    String attrName = randomUTF8();
    String attrValue = randomUTF8();

    AtlasEntityDef classTypeDef = AtlasTypeUtil
            .createClassTypeDef(classType, Collections.<String>emptySet(),
                    AtlasTypeUtil.createUniqueRequiredAttrDef(attrName, "string"));
    AtlasTypesDef atlasTypesDef = new AtlasTypesDef();
    atlasTypesDef.getEntityDefs().add(classTypeDef);
    createType(atlasTypesDef);

    AtlasEntity instance = new AtlasEntity(classType);
    instance.setAttribute(attrName, attrValue);
    AtlasEntityHeader entity = createEntity(instance);
    assertNotNull(entity);
    assertNotNull(entity.getGuid());

    AtlasEntity entityByGuid = getEntityByGuid(entity.getGuid());
    assertEquals(entityByGuid.getAttribute(attrName), attrValue);
}
 

@Test(dependsOnMethods = "testSubmitEntity")
public void testDeleteExistentTraitNonExistentForEntity() throws Exception {

    final String guid = createHiveTable().getGuid();
    final String traitName = "PII_Trait" + randomString();
    AtlasClassificationDef piiTrait = AtlasTypeUtil
            .createTraitTypeDef(traitName, Collections.<String>emptySet(),
                    AtlasTypeUtil.createRequiredAttrDef("type", "string"));
    AtlasTypesDef typesDef = new AtlasTypesDef();
    typesDef.getClassificationDefs().add(piiTrait);
    createType(typesDef);

    try {
        atlasClientV2.deleteClassification(guid, traitName);
        fail("Deletion should've failed for non-existent trait association");
    } catch (AtlasServiceException ex) {
        Assert.assertNotNull(ex.getStatus());
        assertEquals(ex.getStatus(), ClientResponse.Status.BAD_REQUEST);
    }
}
 

@Test(dependsOnMethods = "testSubmitEntity")
public void testGetTraitDefinitionForEntity() throws Exception{
    traitName = "PII_Trait" + randomString();
    AtlasClassificationDef piiTrait =
            AtlasTypeUtil.createTraitTypeDef(traitName, Collections.<String>emptySet());
    AtlasTypesDef typesDef = new AtlasTypesDef();
    typesDef.getClassificationDefs().add(piiTrait);
    createType(typesDef);

    AtlasClassificationDef classificationByName = atlasClientV2.getClassificationDefByName(traitName);
    assertNotNull(classificationByName);

    AtlasEntity hiveTable = createHiveTable();
    assertEquals(hiveTable.getClassifications().size(), 7);

    AtlasClassification piiClassification = new AtlasClassification(piiTrait.getName());

    atlasClientV2.addClassifications(hiveTable.getGuid(), Lists.newArrayList(piiClassification));

    AtlasClassifications classifications = atlasClientV2.getClassifications(hiveTable.getGuid());
    assertNotNull(classifications);
    assertTrue(classifications.getList().size() > 0);
    assertEquals(classifications.getList().size(), 8);
}
 

@BeforeClass
public void setUp() throws Exception {
    new GraphBackedSearchIndexer(typeRegistry);

    // create employee relationship types
    AtlasTypesDef employeeTypes = getDepartmentEmployeeTypes();
    typeDefStore.createTypesDef(employeeTypes);

    AtlasEntitiesWithExtInfo employeeInstances = getDepartmentEmployeeInstances();
    EntityMutationResponse response = entityStore.createOrUpdate(new AtlasEntityStream(employeeInstances), false);

    for (AtlasEntityHeader entityHeader : response.getCreatedEntities()) {
        employeeNameIdMap.put((String) entityHeader.getAttribute(NAME), getAtlasObjectId(entityHeader));
    }

    init();
    AtlasTypesDef typesDef = getInverseReferenceTestTypes();

    AtlasTypesDef typesToCreate = AtlasTypeDefStoreInitializer.getTypesToCreate(typesDef, typeRegistry);

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

public static AtlasTypesDef simpleTypeUpdated(){
    AtlasEntityDef superTypeDefinition =
            createClassTypeDef("h_type", Collections.<String>emptySet(),
                    createOptionalAttrDef("attr", "string"));

    AtlasEntityDef newSuperTypeDefinition =
            createClassTypeDef("new_h_type", Collections.<String>emptySet(),
                    createOptionalAttrDef("attr", "string"));

    AtlasStructDef structTypeDefinition = new AtlasStructDef("s_type", "structType", "1.0",
            Arrays.asList(createRequiredAttrDef("name", "string")));

    AtlasClassificationDef traitTypeDefinition =
            AtlasTypeUtil.createTraitTypeDef("t_type", "traitType", Collections.<String>emptySet());

    AtlasEnumDef enumTypeDefinition = new AtlasEnumDef("e_type", "enumType",
            Arrays.asList(new AtlasEnumElementDef("ONE", "Element Description", 1)));
    AtlasTypesDef ret = AtlasTypeUtil.getTypesDef(Collections.singletonList(enumTypeDefinition), Collections.singletonList(structTypeDefinition),
            Collections.singletonList(traitTypeDefinition), Arrays.asList(superTypeDefinition, newSuperTypeDefinition));

    populateSystemAttributes(ret);

    return ret;
}
 

@Test(dependsOnMethods = "testSubmitEntity")
public void testGetTraitDefinitionForEntity() throws Exception{
    traitName = "PII_Trait" + randomString();
    AtlasClassificationDef piiTrait =
            AtlasTypeUtil.createTraitTypeDef(traitName, ImmutableSet.<String>of());
    AtlasTypesDef typesDef = new AtlasTypesDef();
    typesDef.getClassificationDefs().add(piiTrait);
    createType(typesDef);

    AtlasClassificationDef classificationByName = atlasClientV2.getClassificationDefByName(traitName);
    assertNotNull(classificationByName);

    AtlasEntity hiveTable = createHiveTable();
    assertEquals(hiveTable.getClassifications().size(), 7);

    AtlasClassification piiClassification = new AtlasClassification(piiTrait.getName());

    atlasClientV2.addClassifications(hiveTable.getGuid(), Lists.newArrayList(piiClassification));

    AtlasClassifications classifications = atlasClientV2.getClassifications(hiveTable.getGuid());
    assertNotNull(classifications);
    assertTrue(classifications.getList().size() > 0);
    assertEquals(classifications.getList().size(), 8);
}
 

private void assertNewTypesDef(AtlasTypesDef newTypes, int traitPrayIndex, int vendorPIIIndex, int financeIndex, String expectedTraitPrayStructName, String expectedVendorPIIStructName, String expectedFinanceStructName) {
    assertNotNull(newTypes);
    assertEquals(newTypes.getStructDefs().size(), 4);

    assertTrue(newTypes.getStructDefs().get(traitPrayIndex).getName().contains(LEGACY_TYPE_NAME_PREFIX));
    assertTrue(newTypes.getStructDefs().get(vendorPIIIndex).getName().contains(LEGACY_TYPE_NAME_PREFIX));
    assertTrue(newTypes.getStructDefs().get(financeIndex).getName().contains(LEGACY_TYPE_NAME_PREFIX));

    assertEquals(newTypes.getStructDefs().get(traitPrayIndex).getName(), expectedTraitPrayStructName);
    assertEquals(newTypes.getStructDefs().get(vendorPIIIndex).getName(), expectedVendorPIIStructName);
    assertEquals(newTypes.getStructDefs().get(financeIndex).getName(), expectedFinanceStructName);

    assertEquals(newTypes.getStructDefs().get(1).getAttributeDefs().size(), 1);
    assertEquals(newTypes.getEntityDefs().get(0).getAttributeDefs().get(0).getTypeName(), expectedTraitPrayStructName);
    assertEquals(newTypes.getEntityDefs().get(0).getAttributeDefs().get(1).getTypeName(), String.format("array<%s>", expectedVendorPIIStructName));
    assertEquals(newTypes.getEntityDefs().get(0).getAttributeDefs().get(2).getTypeName(), String.format("map<String,%s>", expectedFinanceStructName));
}
 

@Test
public void performScrub() {
    TypesDefScrubber typesDefScrubber = new TypesDefScrubber();
    AtlasTypesDef td = getTypesDefFromFile(LEGACY_TYPESDEF_JSON);

    int traitPrayIndex = 1;
    int vendorPIIIndex = 2;
    int financeIndex = 3;

    int classificationTraitPrayIndex = 0;
    int classificationVendorPiiIndex = 2;
    int classificationFinancendex = 3;

    String expectedTraitPrayStructName = TypesDefScrubber.getLegacyTypeNameForStructDef(td.getClassificationDefs().get(classificationTraitPrayIndex).getName());
    String expectedVendorPIIStructName = TypesDefScrubber.getLegacyTypeNameForStructDef(td.getClassificationDefs().get(classificationVendorPiiIndex).getName());
    String expectedFinanceStructName = TypesDefScrubber.getLegacyTypeNameForStructDef(td.getClassificationDefs().get(classificationFinancendex).getName());

    assertNewTypesDef(typesDefScrubber.scrub(td), traitPrayIndex, vendorPIIIndex, financeIndex, expectedTraitPrayStructName, expectedVendorPIIStructName, expectedFinanceStructName);

    assertTraitMap(typesDefScrubber, td, classificationTraitPrayIndex, expectedTraitPrayStructName, 0);
    assertTraitMap(typesDefScrubber, td, classificationVendorPiiIndex, expectedVendorPIIStructName, 1);
    assertTraitMap(typesDefScrubber, td, classificationFinancendex, expectedFinanceStructName, 2);
}
 

/**
 * Bulk delete API for all types
 * @param typesDef A composite object that captures all types to be deleted
 * @throws Exception
 * @HTTP 204 On successful deletion of the requested type definitions
 * @HTTP 400 On validation failure for any type definitions
 */
@DELETE
@Path("/typedefs")
@Experimental
public void deleteAtlasTypeDefs(final AtlasTypesDef typesDef) throws AtlasBaseException {
    AtlasPerfTracer perf = null;

    try {
        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "TypesREST.deleteAtlasTypeDefs(" +
                                                           AtlasTypeUtil.toDebugString(typesDef) + ")");
        }



        typeDefStore.deleteTypesDef(typesDef);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

/**
 * Bulk create APIs for all atlas type definitions, only new definitions will be created.
 * Any changes to the existing definitions will be discarded
 * @param typesDef A composite wrapper object with corresponding lists of the type definition
 * @return A composite wrapper object with lists of type definitions that were successfully
 * created
 * @throws Exception
 * @HTTP 200 On successful update of requested type definitions
 * @HTTP 400 On validation failure for any type definitions
 */
@POST
@Path("/typedefs")
public AtlasTypesDef createAtlasTypeDefs(final AtlasTypesDef typesDef) throws AtlasBaseException {
    AtlasPerfTracer perf = null;

    try {
        if (AtlasPerfTracer.isPerfTraceEnabled(PERF_LOG)) {
            perf = AtlasPerfTracer.getPerfTracer(PERF_LOG, "TypesREST.createAtlasTypeDefs(" +
                                                           AtlasTypeUtil.toDebugString(typesDef) + ")");
        }

        return typeDefStore.createTypesDef(typesDef);
    } finally {
        AtlasPerfTracer.log(perf);
    }
}
 

@Test(dependsOnMethods = "testGetTraitNames")
public void testAddTraitWithAttribute() throws Exception {
    final String traitName = "PII_Trait" + randomString();
    AtlasClassificationDef piiTrait = AtlasTypeUtil
            .createTraitTypeDef(traitName, ImmutableSet.<String>of(),
                    AtlasTypeUtil.createRequiredAttrDef("type", "string"));
    AtlasTypesDef typesDef = new AtlasTypesDef();
    typesDef.getClassificationDefs().add(piiTrait);
    createType(typesDef);

    AtlasClassification traitInstance = new AtlasClassification(traitName);
    traitInstance.setAttribute("type", "SSN");

    final String guid = createHiveTable().getGuid();
    atlasClientV2.addClassifications(guid, ImmutableList.of(traitInstance));

    // verify the response
    AtlasEntity withAssociationByGuid = atlasClientV2.getEntityByGuid(guid).getEntity();
    assertNotNull(withAssociationByGuid);
    assertFalse(withAssociationByGuid.getClassifications().isEmpty());

    boolean found = false;
    for (AtlasClassification atlasClassification : withAssociationByGuid.getClassifications()) {
        String attribute = (String)atlasClassification.getAttribute("type");
        if (attribute != null && attribute.equals("SSN")) {
            found = true;
            break;
        }
    }
    assertTrue(found);
}
 

private void updateClassificationDef(AtlasTypesDef typeDefinitionMap, AtlasImportResult result) throws AtlasBaseException {
    for (AtlasClassificationDef def: typeDefinitionMap.getClassificationDefs()) {
        AtlasClassificationDef existing = typeRegistry.getClassificationDefByName(def.getName());
        if(existing != null && addAttributes(existing, def)) {
            typeDefStore.updateClassificationDefByName(existing.getName(), existing);
            result.incrementMeticsCounter("typedef:classification:update");
        }
    }
}
 

private String addNewType() throws Exception {
    String typeName = "test" + randomString();
    AtlasEntityDef classTypeDef = AtlasTypeUtil
            .createClassTypeDef(typeName, ImmutableSet.<String>of(),
                    AtlasTypeUtil.createRequiredAttrDef("name", "string"),
                    AtlasTypeUtil.createRequiredAttrDef("description", "string"));
    AtlasTypesDef typesDef = new AtlasTypesDef();
    typesDef.getEntityDefs().add(classTypeDef);
    createType(typesDef);
    return typeName;
}
 

private List<AtlasClassificationDef> getAllClassificationsDefs() throws Exception {
    MultivaluedMap<String, String> searchParams = new MultivaluedMapImpl();
    searchParams.add(SearchFilter.PARAM_TYPE, "CLASSIFICATION");
    SearchFilter searchFilter = new SearchFilter(searchParams);

    AtlasTypesDef typesDef = atlasClientV2.getAllTypeDefs(searchFilter);
    displayCrLf("Found classifications: " + typesDef.getClassificationDefs().size());
    return typesDef.getClassificationDefs();
}
 

@DataProvider
public Object[][] allCreatedTypes(){
    // Capture all the types that are getting created or updated here.
    AtlasTypesDef updatedTypeDefs = TestUtilsV2.defineValidUpdatedDeptEmployeeTypes();
    AtlasTypesDef allTypeDefs = new AtlasTypesDef();
    allTypeDefs.getEnumDefs().addAll(updatedTypeDefs.getEnumDefs());
    allTypeDefs.getStructDefs().addAll(updatedTypeDefs.getStructDefs());
    allTypeDefs.getClassificationDefs().addAll(updatedTypeDefs.getClassificationDefs());
    allTypeDefs.getEntityDefs().addAll(updatedTypeDefs.getEntityDefs());
    allTypeDefs.getEntityDefs().addAll(TestUtilsV2.getEntityWithValidSuperType());
    return new Object[][] {{allTypeDefs}};
}
 

public void updateTypes(AtlasTypesDef typeDefinitionMap, AtlasImportResult result) throws AtlasBaseException {

        updateEntityDef(typeDefinitionMap, result);
        updateClassificationDef(typeDefinitionMap, result);
        updateEnumDef(typeDefinitionMap, result);
        updateStructDef(typeDefinitionMap, result);
    }
 

public static AtlasTypesDef toAtlasTypesDef(String typeDefinition, AtlasTypeRegistry registry) throws AtlasBaseException {
    AtlasTypesDef ret = new AtlasTypesDef();

    try {
        if (StringUtils.isEmpty(typeDefinition)) {
            throw new AtlasBaseException(INVALID_TYPE_DEFINITION, typeDefinition);
        }

        TypesDef typesDef = AtlasType.fromV1Json(typeDefinition, TypesDef.class);
        if (CollectionUtils.isNotEmpty(typesDef.getEnumTypes())) {
            List<AtlasEnumDef> enumDefs = toAtlasEnumDefs(typesDef.getEnumTypes());
            ret.setEnumDefs(enumDefs);
        }

        if (CollectionUtils.isNotEmpty(typesDef.getStructTypes())) {
            List<AtlasStructDef> structDefs = toAtlasStructDefs(typesDef.getStructTypes());
            ret.setStructDefs(structDefs);
        }

        if (CollectionUtils.isNotEmpty(typesDef.getClassTypes())) {
            List<AtlasEntityDef> entityDefs = toAtlasEntityDefs(typesDef.getClassTypes(), registry);
            ret.setEntityDefs(entityDefs);
        }

        if (CollectionUtils.isNotEmpty(typesDef.getTraitTypes())) {
            List<AtlasClassificationDef> classificationDefs = toAtlasClassificationDefs(typesDef.getTraitTypes());
            ret.setClassificationDefs(classificationDefs);
        }

    } catch (Exception e) {
        LOG.error("Invalid type definition = {}", typeDefinition, e);
        throw new AtlasBaseException(INVALID_TYPE_DEFINITION, typeDefinition);
    }

    return ret;
}
 

private void createBusinessMetadataTypes(String businessMetadataName) throws AtlasBaseException {
    List<AtlasBusinessMetadataDef> businessMetadataDefs = new ArrayList(typesDefs.getBusinessMetadataDefs());
    businessMetadataDefs.add(createBusinessMetadataDef(businessMetadataName));
    typesDefs.setBusinessMetadataDefs(businessMetadataDefs);
    AtlasTypesDef createdTypesDef = typeDefStore.createTypesDef(typesDefs);

    Assert.assertEquals(createdTypesDef.getBusinessMetadataDefs(), businessMetadataDefs, "Data integrity issue while persisting");
}
 

@BeforeTest
public void setup() {
    AtlasTypesDef typesDef = new AtlasTypesDef();
    MockitoAnnotations.initMocks(this);
    when(hiveTableEntityTypeMock.getAttribute("name")).thenReturn(nameAttributeMock);
    when(hiveTableEntityTypeMock.getAttribute("comment")).thenReturn(commentAttributeMock);
    when(hiveTableEntityTypeMock.getAttribute("__state")).thenReturn(stateAttributeMock);
    when(hiveTableEntityTypeMock.getAttribute("description")).thenReturn(descrptionAttributeMock);
    when(hiveTableEntityTypeMock.getAttribute("created")).thenReturn(createdAttributeMock);
    when(hiveTableEntityTypeMock.getAttribute("started")).thenReturn(startedAttributeMock);
    when(hiveTableEntityTypeMock.getAttribute("Constants.ENTITY_TYPE_PROPERTY_KEY")).thenReturn(entitypeAttributeMock);
    when(hiveTableEntityTypeMock.getAttribute("qualifiedName")).thenReturn(qualifiedNameAttributeMock);


    indexFieldNamesMap.put("name", "name_index");
    indexFieldNamesMap.put("comment", "comment_index");
    indexFieldNamesMap.put("__state", "__state_index");
    indexFieldNamesMap.put("description", "descrption__index");
    indexFieldNamesMap.put("created", "created__index");
    indexFieldNamesMap.put("started", "started__index");
    indexFieldNamesMap.put(Constants.ENTITY_TYPE_PROPERTY_KEY, Constants.ENTITY_TYPE_PROPERTY_KEY + "__index");


    when(hiveTableEntityTypeMock.getTypeName()).thenReturn("hive_table");
    when(hiveTableEntityTypeMock2.getTypeName()).thenReturn("hive_db");

    when(nameAttributeMock.getIndexFieldName()).thenReturn("name_index");
    when(commentAttributeMock.getIndexFieldName()).thenReturn("comment_index");
    when(stateAttributeMock.getIndexFieldName()).thenReturn("__state_index");
    when(descrptionAttributeMock.getIndexFieldName()).thenReturn("descrption__index");
    when(createdAttributeMock.getIndexFieldName()).thenReturn("created__index");
    when(startedAttributeMock.getIndexFieldName()).thenReturn("started__index");
    when(entitypeAttributeMock.getIndexFieldName()).thenReturn(Constants.ENTITY_TYPE_PROPERTY_KEY + "__index");
    when(qualifiedNameAttributeMock.getIndexFieldName()).thenReturn("qualifiedName" + "__index");

}
 

@Test(priority = 1)
public void testGet() {
    try {
        AtlasTypesDef typesDef = typeDefStore.searchTypesDef(new SearchFilter());
        assertNotNull(typesDef.getEnumDefs());
        assertEquals(typesDef.getStructDefs().size(), 0);
        assertNotNull(typesDef.getStructDefs());
        assertEquals(typesDef.getClassificationDefs().size(), 0);
        assertNotNull(typesDef.getClassificationDefs());
        assertEquals(typesDef.getEntityDefs().size(), 0);
        assertNotNull(typesDef.getEntityDefs());
    } catch (AtlasBaseException e) {
        fail("Search of types shouldn't have failed");
    }
}
 

public static AtlasTypesDef toAtlasTypesDef(String typeDefinition, AtlasTypeRegistry registry) throws AtlasBaseException {
    AtlasTypesDef ret = new AtlasTypesDef();

    try {
        if (StringUtils.isEmpty(typeDefinition)) {
            throw new AtlasBaseException(INVALID_TYPE_DEFINITION, typeDefinition);
        }

        TypesDef typesDef = TypesSerialization.fromJson(typeDefinition);
        if (CollectionUtils.isNotEmpty(typesDef.enumTypesAsJavaList())) {
            List<AtlasEnumDef> enumDefs = toAtlasEnumDefs(typesDef.enumTypesAsJavaList());
            ret.setEnumDefs(enumDefs);
        }

        if (CollectionUtils.isNotEmpty(typesDef.structTypesAsJavaList())) {
            List<AtlasStructDef> structDefs = toAtlasStructDefs(typesDef.structTypesAsJavaList());
            ret.setStructDefs(structDefs);
        }

        if (CollectionUtils.isNotEmpty(typesDef.classTypesAsJavaList())) {
            List<AtlasEntityDef> entityDefs = toAtlasEntityDefs(typesDef.classTypesAsJavaList(), registry);
            ret.setEntityDefs(entityDefs);
        }

        if (CollectionUtils.isNotEmpty(typesDef.traitTypesAsJavaList())) {
            List<AtlasClassificationDef> classificationDefs = toAtlasClassificationDefs(typesDef.traitTypesAsJavaList());
            ret.setClassificationDefs(classificationDefs);
        }

    } catch (Exception e) {
        LOG.error("Invalid type definition = {}", typeDefinition, e);
        throw new AtlasBaseException(INVALID_TYPE_DEFINITION, typeDefinition);
    }

    return ret;
}
 

public static void loadAllModels(String dirName, AtlasTypeDefStore typeDefStore, AtlasTypeRegistry typeRegistry) throws IOException, AtlasBaseException {
    List<String>  allModels     = getAllModels(dirName);
    for (String model : allModels) {
        AtlasTypesDef typesFromJson = AtlasJson.fromJson(model, AtlasTypesDef.class);
        createTypesAsNeeded(typesFromJson, typeDefStore, typeRegistry);
    }
}
 

private void loadSampleClassificationDefs() throws AtlasBaseException {
    AtlasClassificationDef tag1 = new AtlasClassificationDef("tag1");
    AtlasClassificationDef tag2 = new AtlasClassificationDef("tag2");
    AtlasClassificationDef tag3 = new AtlasClassificationDef("tag3");
    AtlasClassificationDef tag4 = new AtlasClassificationDef("tag4");

    AtlasClassificationDef PII  = new AtlasClassificationDef("PII");
    PII.addAttribute(new AtlasAttributeDef("type", "string"));
    PII.addAttribute(new AtlasAttributeDef("valid", "boolean"));

    typeDefStore.createTypesDef(new AtlasTypesDef(Collections.emptyList(), Collections.emptyList(),
                                                  Arrays.asList(tag1, tag2, tag3, tag4, PII),
                                                  Collections.emptyList(), Collections.emptyList()));
}
 

public static AtlasTypesDef defineInverseReferenceTestTypes() {
    AtlasEntityDef aDef = createClassTypeDef("A", Collections.<String>emptySet(),
        createUniqueRequiredAttrDef("name", "string"),
        new AtlasAttributeDef("b", "B", true, SINGLE, 0, 1, false, false, false, Collections.<AtlasConstraintDef>emptyList()), // 1-1
        new AtlasAttributeDef("oneB", "B", true, SINGLE, 0, 1, false, false, false, Collections.<AtlasConstraintDef>emptyList()), // 1-*
        new AtlasAttributeDef("manyB", AtlasBaseTypeDef.getArrayTypeName("B"), true, SINGLE, 0, 1, false, false, false, Collections.<AtlasConstraintDef>emptyList()),
        new AtlasAttributeDef("mapToB", AtlasBaseTypeDef.getMapTypeName("string", "B"), true, SINGLE, 0, 1, false, false, false,
            Collections.<AtlasConstraintDef>singletonList(new AtlasConstraintDef(
            CONSTRAINT_TYPE_INVERSE_REF, Collections.<String, Object>singletonMap(CONSTRAINT_PARAM_ATTRIBUTE, "mappedFromA"))))); // *-*

    AtlasEntityDef bDef = createClassTypeDef("B", Collections.<String>emptySet(),
        createUniqueRequiredAttrDef("name", "string"),
        new AtlasAttributeDef("a", "A", true, SINGLE, 0, 1, false, false, false,
            Collections.<AtlasConstraintDef>singletonList(new AtlasConstraintDef(
                CONSTRAINT_TYPE_INVERSE_REF, Collections.<String, Object>singletonMap(CONSTRAINT_PARAM_ATTRIBUTE, "b")))),
        new AtlasAttributeDef("manyA", AtlasBaseTypeDef.getArrayTypeName("A"), true, SINGLE, 0, 1, false, false, false,
            Collections.<AtlasConstraintDef>singletonList(new AtlasConstraintDef(
                CONSTRAINT_TYPE_INVERSE_REF, Collections.<String, Object>singletonMap(CONSTRAINT_PARAM_ATTRIBUTE, "oneB")))),
        new AtlasAttributeDef("manyToManyA", AtlasBaseTypeDef.getArrayTypeName("A"), true, SINGLE, 0, 1, false, false, false,
            Collections.<AtlasConstraintDef>singletonList(new AtlasConstraintDef(
                CONSTRAINT_TYPE_INVERSE_REF, Collections.<String, Object>singletonMap(CONSTRAINT_PARAM_ATTRIBUTE, "manyB")))),
        new AtlasAttributeDef("mappedFromA", "A", true, SINGLE, 0, 1, false, false, false, Collections.<AtlasConstraintDef>emptyList()));

    AtlasTypesDef ret = new AtlasTypesDef(Collections.<AtlasEnumDef>emptyList(), Collections.<AtlasStructDef>emptyList(), Collections.<AtlasClassificationDef>emptyList(), Arrays.asList(aDef, bDef));

    populateSystemAttributes(ret);

    return ret;
}