package org.embulk.spi;
import static org.junit.Assert.assertEquals;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import org.embulk.EmbulkTestRuntime;
import org.embulk.config.TaskReport;
import org.embulk.config.ConfigSource;
import org.embulk.config.ConfigDiff;
import org.embulk.config.TaskSource;
import org.embulk.spi.time.Timestamp;
import org.embulk.spi.TestPageBuilderReader.MockPageOutput;
import org.embulk.spi.Schema;
import org.embulk.spi.util.Pages;
public class TestFileInputRunner
{
@Rule
public EmbulkTestRuntime runtime = new EmbulkTestRuntime();
@Before
public void tearDown()
{
MockParserPlugin.raiseException = false;
}
private static class MockFileInputPlugin implements FileInputPlugin
{
Boolean transactionCompleted = null;
Queue<Buffer> buffers;
public MockFileInputPlugin(Queue<Buffer> buffers)
{
this.buffers = buffers;
}
@Override
public ConfigDiff transaction(ConfigSource config,
FileInputPlugin.Control control)
{
control.run(Exec.newTaskSource(), 1);
return null;
}
@Override
public ConfigDiff resume(TaskSource taskSource,
int taskCount,
FileInputPlugin.Control control)
{
throw new UnsupportedOperationException();
}
@Override
public void cleanup(TaskSource taskSource,
int taskCount,
List<TaskReport> successTaskReports)
{
}
public TransactionalFileInput open(TaskSource taskSource,
int taskIndex)
{
return new TransactionalFileInput()
{
@Override
public Buffer poll()
{
return buffers.poll();
}
@Override
public boolean nextFile()
{
return !buffers.isEmpty();
}
@Override
public void close()
{
}
@Override
public void abort()
{
transactionCompleted = false;
}
@Override
public TaskReport commit()
{
transactionCompleted = true;
return null;
}
};
}
}
@Test
public void testMockParserIteration()
{
Buffer[] buffers = new Buffer[] {
runtime.getBufferAllocator().allocate(),
runtime.getBufferAllocator().allocate() };
MockFileInputPlugin fileInputPlugin = new MockFileInputPlugin(
new LinkedList<Buffer>(Arrays.asList(buffers)));
final FileInputRunner runner = new FileInputRunner(fileInputPlugin);
ConfigSource config = Exec.newConfigSource().set(
"parser",
ImmutableMap.of("type", "mock", "columns", ImmutableList.of(
ImmutableMap.of("name", "col1", "type", "boolean", "option", ImmutableMap.of()),
ImmutableMap.of("name", "col2", "type", "long", "option", ImmutableMap.of()),
ImmutableMap.of("name", "col3", "type", "double", "option", ImmutableMap.of()),
ImmutableMap.of("name", "col4", "type", "string", "option", ImmutableMap.of()),
ImmutableMap.of("name", "col5", "type", "timestamp", "option", ImmutableMap.of()),
ImmutableMap.of("name", "col6", "type", "json", "option", ImmutableMap.of()))));
final MockPageOutput output = new MockPageOutput();
runner.transaction(config, new InputPlugin.Control()
{
public List<TaskReport> run(TaskSource inputTaskSource,
Schema schema, int taskCount)
{
List<TaskReport> reports = new ArrayList<>();
reports.add(runner.run(inputTaskSource, schema, 0, output));
return reports;
}
});
assertEquals(true, fileInputPlugin.transactionCompleted);
assertEquals(1, output.pages.size());
Schema schema = config.getNested("parser")
.loadConfig(MockParserPlugin.PluginTask.class)
.getSchemaConfig().toSchema();
List<Object[]> records = Pages.toObjects(schema, output.pages);
assertEquals(2, records.size());
for (Object[] record : records) {
assertEquals(6, record.length);
assertEquals(true, record[0]);
assertEquals(2L, record[1]);
assertEquals(3.0D, (Double) record[2], 0.01D);
assertEquals("45", record[3]);
assertEquals(678L, ((Timestamp) record[4]).toEpochMilli());
assertEquals("{\"_c2\":10,\"_c1\":true,\"_c4\":{\"k\":\"v\"},\"_c3\":\"embulk\"}", record[5].toString());
}
}
@Test
public void testTransactionAborted()
{
Buffer[] buffers = new Buffer[] {
runtime.getBufferAllocator().allocate(),
runtime.getBufferAllocator().allocate() };
MockFileInputPlugin fileInputPlugin = new MockFileInputPlugin(
new LinkedList<Buffer>(Arrays.asList(buffers)));
final FileInputRunner runner = new FileInputRunner(fileInputPlugin);
ConfigSource config = Exec.newConfigSource().set(
"parser",
ImmutableMap.of("type", "mock", "columns", ImmutableList.of(
ImmutableMap.of("name", "col1", "type", "boolean", "option", ImmutableMap.of()),
ImmutableMap.of("name", "col2", "type", "long", "option", ImmutableMap.of()),
ImmutableMap.of("name", "col3", "type", "double", "option", ImmutableMap.of()),
ImmutableMap.of("name", "col4", "type", "string", "option", ImmutableMap.of()),
ImmutableMap.of("name", "col5", "type", "timestamp", "option", ImmutableMap.of()),
ImmutableMap.of("name", "col6", "type", "json", "option", ImmutableMap.of()))));
final MockPageOutput output = new MockPageOutput();
MockParserPlugin.raiseException = true;
try {
runner.transaction(config, new InputPlugin.Control()
{
public List<TaskReport> run(TaskSource inputTaskSource,
Schema schema, int taskCount)
{
List<TaskReport> reports = new ArrayList<>();
reports.add(runner.run(inputTaskSource, schema, 0, output));
return reports;
}
});
} catch (RuntimeException re) {
}
assertEquals(false, fileInputPlugin.transactionCompleted);
assertEquals(0, output.pages.size());
}
}