//Dstl (c) Crown Copyright 2017 package uk.gov.dstl.baleen.consumers.template; import static org.junit.Assert.assertEquals; import java.io.IOException; import java.util.Collection; import java.util.List; import java.util.Map; import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.uima.jcas.tcas.DocumentAnnotation; import org.junit.Before; import uk.gov.dstl.baleen.annotators.testing.AbstractAnnotatorTest; import uk.gov.dstl.baleen.consumers.template.ExtractedRecord.Kind; import uk.gov.dstl.baleen.types.metadata.Metadata; import uk.gov.dstl.baleen.types.templates.TemplateField; import uk.gov.dstl.baleen.types.templates.TemplateRecord; public abstract class AbstractTemplateRecordConsumerTest extends AbstractAnnotatorTest { protected static final String PARA1 = "The quick brown fox jumped over the lazy dog's back."; protected static final String PARA2 = "The quick brown cat jumped over the lazy dog's back."; protected static final String PARA3 = "The quick brown rat jumped over the lazy dog's back."; protected static final String TEXT = String.join("\n", PARA1, PARA2, PARA3); protected final Class<? extends AbstractTemplateRecordConsumer> annotatorClass; protected final String sourceUri; protected final String sourceName; public AbstractTemplateRecordConsumerTest(Class<? extends AbstractTemplateRecordConsumer> clazz) { super(clazz); this.annotatorClass = clazz; this.sourceUri = annotatorClass.getSimpleName() + ".txt"; this.sourceName = annotatorClass.getSimpleName(); } @Before public void beforeAbstractRecordConsumerTest() throws IOException { jCas.setDocumentText(TEXT); DocumentAnnotation documentAnnotation = (DocumentAnnotation) jCas.getDocumentAnnotationFs(); documentAnnotation.setSourceUri(sourceUri); Metadata author = new Metadata(jCas); author.setBegin(0); author.setEnd(0); author.setKey("author"); author.setValue("The Author"); author.addToIndexes(); Metadata creator1 = new Metadata(jCas); creator1.setBegin(0); creator1.setEnd(0); creator1.setKey("creator"); creator1.setValue("The Creator"); creator1.addToIndexes(); Metadata creator2 = new Metadata(jCas); creator2.setBegin(0); creator2.setEnd(0); creator2.setKey("creator"); creator2.setValue("Baleen"); creator2.addToIndexes(); TemplateRecord record1 = new TemplateRecord(jCas); record1.setBegin(0); record1.setEnd(52); record1.setName("record1"); record1.setSource(sourceName); record1.addToIndexes(); TemplateField record1Field1 = new TemplateField(jCas); record1Field1.setBegin(0); record1Field1.setEnd(15); record1Field1.setName("record1Field1"); record1Field1.setSource(sourceName); record1Field1.setValue(record1Field1.getCoveredText()); record1Field1.addToIndexes(); TemplateField record1Field2 = new TemplateField(jCas); record1Field2.setBegin(16); record1Field2.setEnd(31); record1Field2.setName("record1Field2"); record1Field2.setSource(sourceName); record1Field2.setValue(record1Field2.getCoveredText()); record1Field2.addToIndexes(); TemplateRecord record2 = new TemplateRecord(jCas); record2.setBegin(53); record2.setEnd(105); record2.setName("record2"); record2.setSource(sourceName); record2.addToIndexes(); TemplateField record2Field1 = new TemplateField(jCas); record2Field1.setBegin(53); record2Field1.setEnd(68); record2Field1.setName("record2Field1"); record2Field1.setSource(sourceName); record2Field1.setValue(record2Field1.getCoveredText()); record2Field1.addToIndexes(); TemplateField record2Field2 = new TemplateField(jCas); record2Field2.setBegin(69); record2Field2.setEnd(84); record2Field2.setName("record2Field2"); record2Field2.setSource(sourceName); record2Field2.setValue(record2Field2.getCoveredText()); record2Field2.addToIndexes(); TemplateField noRecordField1 = new TemplateField(jCas); noRecordField1.setBegin(106); noRecordField1.setEnd(121); noRecordField1.setName("noRecordField1"); noRecordField1.setSource(sourceName); noRecordField1.setValue(noRecordField1.getCoveredText()); noRecordField1.addToIndexes(); TemplateField noRecordField2 = new TemplateField(jCas); noRecordField2.setBegin(122); noRecordField2.setEnd(137); noRecordField2.setName("noRecordField2"); noRecordField2.setSource(sourceName); noRecordField2.setValue(noRecordField2.getCoveredText()); noRecordField2.addToIndexes(); } protected void checkRecords(Map<String, Collection<ExtractedRecord>> recordMap) { Collection<ExtractedRecord> records = recordMap.get(annotatorClass.getSimpleName()); Stream<ExtractedRecord> recordStream = records.stream() .filter(p -> p.getKind().equals(Kind.NAMED) && p.getName().equals("record1")); List<ExtractedRecord> collect = recordStream.collect(Collectors.toList()); ExtractedRecord record1 = collect.get(0); assertEquals(Kind.NAMED, record1.getKind()); assertEquals(2, record1.getFields().size()); assertEquals("The quick brown", findFieldValue("record1Field1", record1.getFields())); assertEquals("fox jumped over", findFieldValue("record1Field2", record1.getFields())); ExtractedRecord record2 = records.stream() .filter(p -> p.getKind().equals(Kind.NAMED) && p.getName().equals("record2")) .collect(Collectors.toList()).get(0); assertEquals(Kind.NAMED, record1.getKind()); assertEquals(2, record2.getFields().size()); assertEquals("The quick brown", findFieldValue("record2Field1", record2.getFields())); assertEquals("cat jumped over", findFieldValue("record2Field2", record2.getFields())); ExtractedRecord defaultRecord = records.stream().filter(p -> p.getKind().equals(Kind.DEFAULT)) .collect(Collectors.toList()).get(0); assertEquals(null, defaultRecord.getName()); assertEquals(2, defaultRecord.getFields().size()); assertEquals("The quick brown", findFieldValue("noRecordField1", defaultRecord.getFields())); assertEquals("rat jumped over", findFieldValue("noRecordField2", defaultRecord.getFields())); } protected String findFieldValue(String fieldName, Collection<ExtractedField> fields) { for (ExtractedField extractedField : fields) { if (extractedField.getName().equals(fieldName)) { return extractedField.getValue(); } } return null; } }