NotificationEntityChangeListener implements EntityChangeListener { @VisibleForTesting public static List<IStruct> getAllTraits(IReferenceableInstance entityDefinition, TypeSystem typeSystem) throws AtlasException { List<IStruct> traitInfo = new LinkedList<>(); for (String traitName : entityDefinition.getTraits()) { IStruct trait = entityDefinition.getTrait(traitName); String typeName = trait.getTypeName(); Map<String, Object> valuesMap = trait.getValuesMap(); traitInfo.add(new Struct(typeName, valuesMap)); traitInfo.addAll(getSuperTraits(typeName, valuesMap, typeSystem)); } return traitInfo; } @Inject  NotificationEntityChangeListener(NotificationInterface notificationInterface, TypeSystem typeSystem); @Override void onEntitiesAdded(Collection<ITypedReferenceableInstance> entities, boolean isImport); @Override void onEntitiesUpdated(Collection<ITypedReferenceableInstance> entities, boolean isImport); @Override void onTraitsAdded(ITypedReferenceableInstance entity, Collection<? extends IStruct> traits); @Override void onTraitsDeleted(ITypedReferenceableInstance entity, Collection<String> traitNames); @Override void onTraitsUpdated(ITypedReferenceableInstance entity, Collection<? extends IStruct> traits); @Override void onEntitiesDeleted(Collection<ITypedReferenceableInstance> entities, boolean isImport); @VisibleForTesting static List<IStruct> getAllTraits(IReferenceableInstance entityDefinition,
                                              TypeSystem typeSystem);  }
@Test public void testGetAllTraitsSuperTraits() throws Exception { TypeSystem typeSystem = mock(TypeSystem.class); String traitName = "MyTrait"; IStruct myTrait = new Struct(traitName); String superTraitName = "MySuperTrait"; TraitType traitDef = mock(TraitType.class); Set<String> superTypeNames = Collections.singleton(superTraitName); TraitType superTraitDef = mock(TraitType.class); Set<String> superSuperTypeNames = Collections.emptySet(); Referenceable entity = getEntity("id", myTrait); when(typeSystem.getDataType(TraitType.class, traitName)).thenReturn(traitDef); when(typeSystem.getDataType(TraitType.class, superTraitName)).thenReturn(superTraitDef); when(traitDef.getAllSuperTypeNames()).thenReturn(superTypeNames); when(superTraitDef.getAllSuperTypeNames()).thenReturn(superSuperTypeNames); List<IStruct> allTraits = NotificationEntityChangeListener.getAllTraits(entity, typeSystem); assertEquals(2, allTraits.size()); for (IStruct trait : allTraits) { String typeName = trait.getTypeName(); assertTrue(typeName.equals(traitName) || typeName.equals(superTraitName)); } }