//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;
}
}