embulk java-filter plugin のテストの書き方

こんにちは。@Civitaspo です。先日12/15 の Embulk Meetup Tokyo#2 では、『Embulkに足りない5つのコト』という喧嘩腰なタイトルで発表させていただきました。おかげで色んな方とembulkについて議論できた良いMeetupになりました。呼んでくださった @frsyuki と、裏から支えて下さった @sonots にはこの場をお借りして感謝の意を伝えたいと思います。ありがとうございました。

さて、今日は embulk java-filter plugin のテストの書き方 について記事を書こうと思います。

おことわり

• 話を簡単にするためにjava-filter pluginのテストの書き方にfocusしていますが、他のtypeのjava pluginでも応用できる内容だと思っています。
• EmbulkEmbedを使ったテストも書けますが今回は紹介しません。
• 依存のないClassやMethodレベルのテストは普通に書けるだろうと思うので言及はしません。
• 「こんな書き方の方が良いぜよ」「ここ分かりにくい」など、ご意見ありましたらドシドシ連絡いただきたいです！
  • (書いた後に思いましたがあまり一般化されていないかもしれません。。。)

参考資料

• Embulk Plugins
• TestRenameFilterPlugin.java
• embulk-coreのtestライブラリ
  • EmbulkTestRuntime
• Embulk Javaパーサープラグインの単体テスト
  • EmbulkEmbedを使ったテストの書き方を紹介しています
• JUnit Rules
• JUnit4.7 の新機能 Rules とは
• 並列バッチデータ転送OSSのEmbulkをソースコードリーディングしてみる（その３：run概要＆データソース概要
  • EmbulkTestRuntime を理解するために読んでおくと吉
• ExpectedException ルールを使いこなしたい

はじめに

filter pluginのテストとしては

• configを正しく処理できているか
• inputのデータを処理した後、outputのデータが意図した状態になっているか

がわかれば良いと思っているので上記2点に関するテストの書き方を説明します。
説明の都合上、拙作の embulk-filter-expand_json のコードをサンプルとして説明していきます。

準備編

依存関係の追加

build.gradle に以下のように依存関係を追加します。

// ...

dependencies {
    compile  "org.embulk:embulk-core:0.8.+"
    provided "org.embulk:embulk-core:0.8.+"
    // compile "YOUR_JAR_DEPENDENCY_GROUP:YOUR_JAR_DEPENDENCY_MODULE:YOUR_JAR_DEPENDENCY_VERSION"
    testCompile "junit:junit:4.+"
    testCompile "org.embulk:embulk-core:0.8.+:tests" // <= コレ!!
}

// ...

これで embulk-coreのtestライブラリが使えるようになります。

EmbulkTestRuntime を Rule に設定する

JUnitにはRulesという機能があります。EmbulkのテストライブラリにはこのRuleに使用できる EmbulkTestRuntime という Classが実装されています。

// ...
import org.embulk.EmbulkTestRuntime;
import org.junit.Rule;
// ...

public class TestExpandJsonFilterPlugin
{
    @Rule
    public EmbulkTestRuntime runtime = new EmbulkTestRuntime();

    // ...
}

このClassをRuleに設定することで、EmbulkのRuntime時でないと行えないテストが行えるようになります。

YAMLのconfigをロードするメソッドを作っておく

実際に設定したYAMLをEmbulkにロードしてテスト結果を取得するためにYAMLのconfigをロードするメソッドを定義しておきます。

// ...
import org.embulk.config.ConfigLoader;
import org.embulk.config.ConfigSource;
// ...

public class TestExpandJsonFilterPlugin
{
    // ...

    private ConfigSource getConfigFromYaml(String yaml)
    {
        ConfigLoader loader = new ConfigLoader(Exec.getModelManager());
        return loader.fromYamlString(yaml);
    }

    // ...
}

これでYAMLのconfigをPluginにロードする準備ができました。

実際にテストを書く編

configを正しく処理できているか

configのテストは以下の2点を見ておけば良いと思います。

• requiredなparamに値がセットされない場合にexceptionを吐くか
• optionalなparamに値がセットされない場合に定義したdefaultの値が入っているか

requiredなparamに値がセットされない場合にexceptionを吐くかのテストを書く

準備

exceptionをtestするために、Rules に ExpectedException を設定しておきます。

// ...
import org.junit.rules.ExpectedException;
// ...

public class TestExpandJsonFilterPlugin
{
    // ...

    @Rule
    public ExpectedException exception = ExpectedException.none();

    // ...
}

テストを書く

// ...
import org.embulk.config.ConfigException;
import org.junit.Test;
import static org.embulk.filter.expand_json.ExpandJsonFilterPlugin.PluginTask;
// ...

public class TestExpandJsonFilterPlugin
{
    // ...

    // json_column_name という required param が設定されていない場合に
    // exception を履くことを確認するテスト
    @Test
    public void testThrowExceptionAbsentJsonColumnName()
    {
        String configYaml = "" +
                "type: expand_json\n" +
                "expanded_columns:\n" +
                "  - {name: _c1, type: string}";
        ConfigSource config = getConfigFromYaml(configYaml);

        exception.expect(ConfigException.class);
        exception.expectMessage("Field 'json_column_name' is required but not set");
        config.loadConfig(PluginTask.class);
    }

    // ...
}

String configYaml にfilter pluginに設定するconfigをセットし、先ほど準備した getConfigFromYaml に食わせて、 ConfigSource を取得しておきます。これ以降この設定がよく出てきますが逐一、言及しないことにします。

次に exception.expect と exception.expectMessage で発生する例外について定義します。 configの例外は org.embulk.config.ConfigException なので、設定します。

最後に config を plugin に読み込ませます。 config.loadConfig(PluginTask.class); は transaction の一番初めに呼ばれるのでお馴染みですね。この config のロードで例外が発生します。

optionalなparamに値がセットされない場合に定義したdefaultの値が入っているかのテストを書く

次は config をロードした結果、getterで取得した値をassertしましょう。

public class TestExpandJsonFilterPlugin
{
    // ...

    @Test
    public void testDefaultValue()
    {
        String configYaml = "" +
                "type: expand_json\n" +
                "json_column_name: _c0\n" +
                "expanded_columns:\n" +
                "  - {name: _j1, type: boolean}\n" +
                "  - {name: _j2, type: long}\n" +
                "  - {name: _j3, type: timestamp}\n" +
                "  - {name: _j4, type: double}\n" +
                "  - {name: _j5, type: string}\n";

        ConfigSource config = getConfigFromYaml(configYaml);
        PluginTask task = config.loadConfig(PluginTask.class);

        assertEquals("$.", task.getRoot());
        assertEquals("UTC", task.getTimeZone());
        assertEquals("%Y-%m-%d %H:%M:%S.%N %z", task.getDefaultTimestampFormat());
    }

    // ...
}

PluginTask task = config.loadConfig(PluginTask.class); で PluginTask を取得し、定義されたgetterを利用します。

inputのデータを処理した後、outputのデータが意図した状態になっているか

inputのデータを処理した後、outputのデータが意図した状態になっているかのテストは以下の2点を見ておけば良いと思います。

• input schema / output schema が正しく変更されているか
• input data / output data が正しく変更されているか

input schema / output schema が正しく変更されているかのテストを書く

transaction を呼び出してschemaの変更をテストします。

// ...
import org.embulk.spi.Schema;
import org.embulk.spi.Column;
import static org.embulk.filter.expand_json.ExpandJsonFilterPlugin.Control;
// ...

public class TestExpandJsonFilterPlugin
{
    // ...

    @Test
    public void testExpandJsonKeyToSchema()
    {
        String configYaml = "" +
                "type: expand_json\n" +
                "json_column_name: _c0\n" +
                "root: $.\n" +
                "expanded_columns:\n" +
                "  - {name: _j1, type: boolean}\n" +
                // ...
                "  - {name: _c0, type: string}\n";

        ConfigSource config = getConfigFromYaml(configYaml);

        final Schema inputSchema = Schema.builder()
            .add("_c0", STRING)
            .add("_c1", STRING)
            .build();

        ExpandJsonFilterPlugin expandJsonFilterPlugin = new ExpandJsonFilterPlugin();
        expandJsonFilterPlugin.transaction(config, inputSchema, new Control()
        {
            @Override
            public void run(TaskSource taskSource, Schema outputSchema)
            {
                assertEquals(7, outputSchema.getColumnCount());

                Column new_j1 = outputSchema.getColumn(0);
                // ...

                assertEquals("_j1", new_j1.getName());
                assertEquals(BOOLEAN, new_j1.getType());
                // ...
            }
        });
    }

    // ...
}

        final Schema inputSchema = Schema.builder()
            .add("_c0", STRING)
            .add("_c1", STRING)
            .build();

        ExpandJsonFilterPlugin expandJsonFilterPlugin = new ExpandJsonFilterPlugin();
        expandJsonFilterPlugin.transaction(config, inputSchema, new Control()

ここで plugin classのインスタンスを生成し、 transaction メソッドを呼び出します。第三引数には org.embulk.filter.expand_json.ExpandJsonFilterPlugin.Control のインタフェースを渡して、その中で処理結果のテストを行います。

            @Override
            public void run(TaskSource taskSource, Schema outputSchema)

run メソッドの中に outputSchema が渡ってくるので、このschemaをassertします。

input data / output data が正しく変更されているかのテストを書く

// ...
import org.embulk.spi.Exec;
import org.embulk.spi.Page;
import org.embulk.spi.PageOutput;
import org.embulk.spi.PageReader;
import org.embulk.spi.PageTestUtils;
import org.embulk.spi.Schema;
import org.embulk.spi.TestPageBuilderReader.MockPageOutput;
// ...

public class TestExpandJsonFilterPlugin
{
    // ...

    @Test
    public void testExpandJsonValues()
    {
        String configYaml = "" +
                "type: expand_json\n" +
                "json_column_name: _c0\n" +
                "root: $.\n" +
                "time_zone: Asia/Tokyo\n" +
                "expanded_columns:\n" +
                "  - {name: _j0, type: boolean}\n" +
                // ...
                "  - {name: _c0, type: string}\n";

        ConfigSource config = getConfigFromYaml(configYaml);

        final Schema inputSchema = Schema.builder()
            .add("_c0", STRING)
            .add("_c1", STRING)
            .build();

        ExpandJsonFilterPlugin expandJsonFilterPlugin = new ExpandJsonFilterPlugin();
        expandJsonFilterPlugin.transaction(config, inputSchema, new Control()
        {
            @Override
            public void run(TaskSource taskSource, Schema outputSchema)
            {
                MockPageOutput mockPageOutput = new MockPageOutput();
                PageOutput pageOutput = expandJsonFilterPlugin.open(taskSource,
                                                                    inputSchema,
                                                                    outputSchema,
                                                                    mockPageOutput);

                ImmutableMap.Builder<String,Object> builder = ImmutableMap.builder();
                builder.put("_j0", true);
                String data = convertToJsonString(builder.build()); // MapをJSON文字列に変換している

                for (Page page : PageTestUtils.buildPage(runtime.getBufferAllocator(),
                                                         inputSchema,
                                                         data, "_c1_data")) {
                    pageOutput.add(page);
                }

                pageOutput.finish();
                pageOutput.close();

                PageReader pageReader = new PageReader(outputSchema);

                for (Page page : mockPageOutput.pages) {
                    pageReader.setPage(page);
                    assertEquals(true, pageReader.getBoolean(outputSchema.getColumn(0)));
                    // ...
                    assertEquals(c1Data,
                                 pageReader.getString(outputSchema.getColumn(13)));
                }
            }
        });
    }


    // ...
}

先程と同様、 org.embulk.filter.expand_json.ExpandJsonFilterPlugin.Control の run メソッドの中でテストを行ってきます。

                MockPageOutput mockPageOutput = new MockPageOutput();
                // ...
                for (Page page : mockPageOutput.pages) {

MockPageOutput を使います。これは名前の通り Mock化した PageOutput を提供してくれます。この MockPageOutput の pages というメソッドから Page を取得することができます。

                PageOutput pageOutput = expandJsonFilterPlugin.open(taskSource,
                                                                    inputSchema,
                                                                    outputSchema,
                                                                    mockPageOutput);

その MockPageOutput を使ってPluginの open メソッドを呼んで PageOutput を生成します。

                for (Page page : PageTestUtils.buildPage(runtime.getBufferAllocator(),
                                                         inputSchema,
                                                         data, "_c1_data")) {
                    pageOutput.add(page);
                }

テスト用のデータを先ほど生成した pageOutput に add していきます。 テスト用のデータ生成には PageTestUtils の buildPage というメソッドを使用します。第一引数には EmbulkTestRuntime から取得した BufferAllocator を、第二引数には inputSchema を、 第三引数以降は実際のデータを渡します。

                PageReader pageReader = new PageReader(outputSchema);

                for (Page page : mockPageOutput.pages) {
                    pageReader.setPage(page);
                    assertEquals(true, pageReader.getBoolean(outputSchema.getColumn(0)));
                    // ...
                    assertEquals(c1Data,
                                 pageReader.getString(outputSchema.getColumn(13)));
                }

後は実際に取得できるデータをassertしていきます。PageReader を使って取得します。

さいごに

なんだか僕が書いたコードのコードリーディングみたいな説明になってしまいましたがfilter pluginのテストの書き方はご理解いただけたでしょうか。
Embulk Pluginsはテストがないものが多いのでこの機会にテストを書いて見ようとなる人がいれば幸いです。