less is more

心のフルスタックエンジニア👨‍💻バイブスでコードを書いています🤘

cobraでGoのCLIツールを作る

概要

この記事で作ったCLIツールの裏側の話です。

bluepixel.hatenablog.com

OSSとしてバイナリの配布も始めました。

github.com

GoでCLIツールを作る手段

f:id:bluepixel:20200506182434p:plain いくつか候補が見つかったんですが、何分 Go に入門したばかりなのでPros/Consが判断できず。

とりあえず Kuberntes や Hugo など大手プロジェクトに採用実績があり、viper での設定の注入など機能が豊富そうな cobra で雛形を作ってみることにしました。

github.com

サブコマンドで実装

cobraの使い方自体は日本語記事もいくつかあるので特に解説はしません。
今回作るコマンドはサブコマンドとして実装します。

このABCというライブラリには将来的に別の機能をもったサブコマンドをどんどん生やしていく予定です。

cobraでは以下のコマンドを実行するとサブコマンドの雛形が生成されます。

cobra add サブコマンド名
$ cobra add sub_command
subCommand created at /xxxx/xxx/xxx
$ ls cmd
root.go     subCommand.go
// subCommand.go
package cmd

import (
    "fmt"

    "github.com/spf13/cobra"
)

var subCommandCmd = &cobra.Command{
    Use:   "subCommand",
    Short: "A brief description of your command",
    Long: `A longer description that spans multiple lines and likely contains examples
and usage of using your command. For example:

Cobra is a CLI library for Go that empowers applications.
This application is a tool to generate the needed files
to quickly create a Cobra application.`,
    Run: func(cmd *cobra.Command, args []string) {
        fmt.Println("subCommand called")
    },
}

func init() {
    rootCmd.AddCommand(subCommandCmd)
}

Runに中身を実装していくわけですが、その前に、このままだと扱いづらいので、この雛形自体を作り直します。

具体的には、ユニットテストがしづらい、入出力のストリームを握りづらい、エントリーポイントは薄く保ちたい、などです。

雛形の改修

まずはコマンドの初期化をはがします。 lib/に新しくパケージを切って、そのコマンドに関する定義は全てそこに閉じ込めます。

そして新たに*cobra.Commandを返す関数を作って以下を置き換えます。
そうすることでテストなどで取り回しやすくなります。

var subCommandCmd = &cobra.Command{
    Use:   "subCommand",
    ...
// lib/sub/sub.go
package sub

import (
    "github.com/spf13/cobra"
)

func NewCmd() *cobra.Command {
    cmd := &cobra.Command{
        Use:   "sub",
        Short: "this is sub command",
        Long:  `this is sub command. this is sub command. this is sub command. this is sub command. this is sub command. this is sub command.`,
        RunE: func(cmd *cobra.Command, args []string) error {
            cmd.Println("hello from sub command")
            return nil
        },
    }
    return cmd
}

エントリーポイントの方はこう。

// cmd/sub.go
package cmd

import (
    "github.com/Blue-Pix/abc/lib/sub"
)

var subCmd = sub.NewCmd()

func init() {
    rootCmd.AddCommand(subCmd)
}

次にサブコマンドからルートコマンドへ出力を流すようにします。
これは、テストで出力ストリームを見る際に、ルートコマンド経由でサブコマンドをテストできるようにするためです。

OutOrStdout()で現在設定されている出力ストリーム(デフォルトでは標準出力)が取れる。

func init() {
    subCmd.SetOut(rootCmd.OutOrStdout())
    rootCmd.AddCommand(subCmd)
}

一つ注意しなければいけないのは、コマンドの処理内ではcmd.Printlnのようにcmd.*を使うこと。
いつも通りカジュアルにfmt.Printlnにしているとコマンドに設定した出力ストリームに流れない。

こんな感じの雛形の変更をデフォルトのrootCmdに対しても行います。

中身の実装

エントリーポイントから分離したパッケージに内部実装を書いていきます。

ここからは実際に実装したabc amiコマンドの例で見ていきます。 github.com

まずはフラグの取り方から。
NewCmd()の初期化処理に挟み込みます。

複数のコマンド間で共通で取れるPersistentFlagsと当該コマンド内のみで有効なLocalFlagsの2種類があるのですが、今回はLocalFlagsにする。

func NewCmd() *cobra.Command {
    cmd := &cobra.Command{
        ....
    }
    cmd.Flags().StringP("version", "v", "", "os version(1 or 2)")
    cmd.Flags().StringP("virtualization-type", "V", "", "virtualization type(hvm or pv)")
    cmd.Flags().StringP("arch", "a", "", "cpu architecture(x86_64 or arm64)")
    cmd.Flags().StringP("storage", "s", "", "storage type(gp2, ebs or s3)")
    cmd.Flags().StringP("minimal", "m", "", "if minimal image or not(true or false)")
    return cmd
}

intやbooleanにできるものもあるのだが、JSON文字列で扱ったりもするので、stringの方が都合が良いのでこうした。

受ける側ではcmd.Flags().GetString("version")のようにしてコマンドラインから与えられた引数を受け取り、フィルタリングの関数を呼んでいく。

amis := getAMIList()
if version, err := cmd.Flags().GetString("version"); version != "" && err == nil {
    amis = filterByVersion(version, amis)
}
if virtualizationType, err := cmd.Flags().GetString("virtualization-type"); virtualizationType != "" && err == nil {
    amis = filterByVirtualizationType(virtualizationType, amis)
}
if arch, err := cmd.Flags().GetString("arch"); arch != "" && err == nil {
    amis = filterByArch(arch, amis)
}

そしてこちらでもテスタビリティのために、出力処理を境界に分離する。

func run(cmd *cobra.Command, args []string) {
    str := Run(cmd, args)
    cmd.Println(str)
}

func Run(cmd *cobra.Command, args []string) string {
    amis := getAMIList()

    if version, err := cmd.Flags().GetString("version"); version != "" && err == nil {
        amis = filterByVersion(version, amis)
    }
    if virtualizationType, err := cmd.Flags().GetString("virtualization-type"); virtualizationType != "" && err == nil {
        amis = filterByVirtualizationType(virtualizationType, amis)
    }
    if arch, err := cmd.Flags().GetString("arch"); arch != "" && err == nil {
        amis = filterByArch(arch, amis)
    }
    if storage, err := cmd.Flags().GetString("storage"); storage != "" && err == nil {
        amis = filterByStorage(storage, amis)
    }
    if minimal, err := cmd.Flags().GetString("minimal"); minimal != "" && err == nil {
        amis = filterByMinimal(minimal, amis)
    }
    str := toJSON(amis)
    return str
}

テストを書く

当初はサブコマンドごとにテストを書こうとしていたのだが、
argsとして渡した引数がフラグとしてパースされずにargsに残ってしまう問題があり、
普通にルートコマンド経由でサブコマンドの文字列も含めて渡してテストするようにした。

assertは以下のように出力ストリームを奪って比較する。

b := bytes.NewBufferString("")
cmd.SetOut(b)
cmd.Execute()
out, err := ioutil.ReadAll(b)
if err != nil {
    t.Fatal(err)
}
// if out == ???

実際のテストコード。
出力されたJSON文字列の正規表現のマッチを見る。
通常のオプション、短縮オプションを両方試す。

// lib/root/root_test.go
package root

func prepareCmd(args []string) *cobra.Command {
    cmd := NewCmd()
    cmd.SetArgs(args)
    amiCmd := ami.NewCmd()
    cmd.AddCommand(amiCmd)
    return cmd
}

func TestExecute(t *testing.T) {
    t.Run("ami", func(t *testing.T) {
        t.Run("query by --version", func(t *testing.T) {
            args := []string{"ami", "--version", "2"}
            cmd := prepareCmd(args)
            b := bytes.NewBufferString("")
            cmd.SetOut(b)
            cmd.Execute()
            out, err := ioutil.ReadAll(b)
            if err != nil {
                t.Fatal(err)
            }
            r := regexp.MustCompile("\"version\":\"([^\"]+)\"")
            list := r.FindAllStringSubmatch(string(out), -1)
            if len(list) == 0 {
                t.Fatal("there is no much result")
            }
            for _, l := range list {
                if l[1] != args[2] {
                    t.Fatal(fmt.Sprintf("expected: %s, actual: %s", args[2], l[1]))
                }
            }
        })

        t.Run("query by -v", func(t *testing.T) {
            args := []string{"ami", "-v", "1"}
            cmd := prepareCmd(args)
            b := bytes.NewBufferString("")
            cmd.SetOut(b)
            cmd.Execute()
            out, err := ioutil.ReadAll(b)
            if err != nil {
                t.Fatal(err)
            }
            r := regexp.MustCompile("\"version\":\"([^\"]+)\"")
            list := r.FindAllStringSubmatch(string(out), -1)
            if len(list) == 0 {
                t.Fatal("there is no much result")
            }
            for _, l := range list {
                if l[1] != args[2] {
                    t.Fatal(fmt.Sprintf("expected: %s, actual: %s", args[2], l[1]))
                }
            }
        })
    })
}

goの標準パッケージの正規表現はre2で実装されていて否定的先読みが使えなかったので、愚直にマッチした内容を見ている。

正規表現一発でいける場合はassert.Regexpが使える。
否定はassert.NotRegexpで。
NotRegexpって。

所感

CLIツールのテストどこまでやるか問題の課題。
今回書いたテストは結構E2Eまで踏み込んでいるが、
ユニットテストだけで終わらせるとなかなかCLIとのインターフェースのつなぎこみ部分でチェックできないところが多く出てくるので難しいところ。

あとは外部APIとのモックに関して。
今回はRead系の操作なのでモックせずに実際にAPIを叩いている。
AWSの認証まわりは実際に動かさないと確認できないし、
ツールとして根幹の機能が提供できていることを担保するためにもこのあたりはなるべく動かして見るべきだと思っている。
Create系はリソース後片付けがあるのでなかなか難しいが。

cobraに関しては、良いかどうかよくわからない。
viper使ったり、共通のフラグがある場合は、そのあたりの透過的な扱いの恩恵を受けられると思うが、構成によっては自前でflagなりpflagなりを管理した方がやりやすいのかもしれない。

もう2,3コマンド実装してみよう。