EntityExtractionBolt extends ChatAlyticsBaseBolt { @VisibleForTesting protected List<ChatEntity> extractEntities(FatMessage fatMessage) { Message message = fatMessage.getMessage(); String messageStr = message.getMessage(); if (messageStr == null) { return ImmutableList.of(); } List<Triple<String,Integer,Integer>> classification = classifier.classifyToCharacterOffsets(messageStr); Map<String, ChatEntity> entities = Maps.newHashMapWithExpectedSize(classification.size()); for (Triple<String, Integer, Integer> triple : classification) { if (triple.third - triple.second > MAX_ENTITY_CHARS) { continue; } String entity = messageStr.substring(triple.second, triple.third); ChatEntity existingEntity = entities.remove(entity); int occurrences; if (existingEntity == null) { occurrences = 1; } else { occurrences = existingEntity.getOccurrences() + 1; } Room room = fatMessage.getRoom(); String roomName = null; if (room != null) { roomName = room.getName(); } entities.put(entity, new ChatEntity(fatMessage.getUser().getMentionName(), roomName, message.getDate(), entity, occurrences, fatMessage.getUser().isBot())); } LOG.debug("Extracted {} entities", entities.size()); return Lists.newArrayList(entities.values()); }  @Override void prepare(ChatAlyticsConfig config, @SuppressWarnings("rawtypes") Map conf,
                        TopologyContext context); @Override void execute(Tuple input, BasicOutputCollector collector); @Override void declareOutputFields(OutputFieldsDeclarer fields); @Override void cleanup(); static final String BOLT_ID; static final String CHAT_ENTITY_FIELD_STR; }
@Test public void testExtractEntities() { DateTime date = DateTime.now().withZone(DateTimeZone.UTC); String mentionName = "jane"; String userId = "1"; String roomId = "100"; String ent1 = "Jane Doe"; String ent2 = "Mount Everest"; Message msg = new Message(date, mentionName, userId, String.format("Today, %s is going to climb %s", ent1, ent2), roomId, MessageType.MESSAGE); FatMessage fatMessage = new FatMessage(msg, user, room); List<ChatEntity> entities = underTest.extractEntities(fatMessage); Map<String, ChatEntity> entitiesMap = Maps.newHashMapWithExpectedSize(entities.size()); entities.forEach((entity) -> entitiesMap.put(entity.getValue(), entity)); assertEquals(2, entities.size()); ChatEntity entity = entitiesMap.get(ent1); assertEquals(ent1, entity.getValue()); assertEquals(1, entity.getOccurrences()); assertEquals(date, entity.getMentionTime()); entity = entitiesMap.get(ent2); assertEquals(ent2, entity.getValue()); assertEquals(1, entity.getOccurrences()); assertEquals(date, entity.getMentionTime()); }
@Test public void testExtractEntities_nullMessage() { DateTime date = DateTime.now().withZone(DateTimeZone.UTC); Message msg = new Message(date, "jane", "1", null, "100", MessageType.MESSAGE); User mockUser = mock(User.class); when(mockUser.getMentionName()).thenReturn("jane"); Room mockRoom = mock(Room.class); when(mockRoom.getName()).thenReturn("theroom"); FatMessage fatMessage = new FatMessage(msg, mockUser, mockRoom); List<ChatEntity> result = underTest.extractEntities(fatMessage); assertTrue(result.isEmpty()); }
@Test public void testExtractEntities_doubleOccurrence() { String mentionName = "jane"; String userId = "1"; String roomId = "100"; String ent = "Mount Everest"; Message msg = new Message(DateTime.now(), mentionName, userId, String.format("Today, I'm going to climb %s and %s", ent, ent), roomId, MessageType.MESSAGE); FatMessage fatMessage = new FatMessage(msg, user, room); List<ChatEntity> entities = underTest.extractEntities(fatMessage); assertEquals(1, entities.size()); assertEquals(2, entities.get(0).getOccurrences()); }
@Test public void testExtractEntities_nullRoom() { String mentionName = "jane"; String userId = "1"; String roomId = "100"; String ent = "Mount Everest"; Message msg = new Message(DateTime.now(), mentionName, userId, String.format("Today, I'm going to climb %s", ent), roomId, MessageType.MESSAGE); FatMessage fatMessage = new FatMessage(msg, user, null); List<ChatEntity> entities = underTest.extractEntities(fatMessage); assertEquals(1, entities.size()); assertNull(entities.get(0).getRoomName()); }