Go言語での集合(Set)の扱い方とテスト

はじめに

Go言語には標準で集合(Set)は用意されていません。

実現方法

1. Slice

strSet1 := []string{"element1", "element2", "element3"}

シンプルな方法。
重複は許されるし、ただのスライス。
簡易的な場合はこれで十分。

2. map[interface{}]struct{}

strSet2 := map[string]struct{}{
	"element1": struct{}{},
	"element2": struct{}{},
	"element3": struct{}{},
}

mapを利用する方法。
struct{}はメモリを圧迫することはない。

3. github.com/deckarep/golang-set/mapset

strSet3 := mapset.NewSet()
requiredClasses.Add("element1")
requiredClasses.Add("element2")
requiredClasses.Add("element3")

パッケージを利用する方法。
Setを操作する便利メソッドが多く実装されている。
内部実装ではmap[interface{}]struct{}が用いられている。
https://github.com/deckarep/golang-set にソースコードがある。

用途

個人開発しているときに、RedisのSMembersのテストをする必要があった。
RedisではSetでデータを保持することができ、SMembersで取得することができる。
goにはSet型はないため、go-redisのSMembersでは[]stringが返る。
集合どうしを比較する場合に[]stringでは要素間の順序が担保されず、テストが失敗するためSet型を利用した。
3のgithub.com/deckarep/golang-set/mapset を利用するのは大袈裟と思い、2のmap[interface{}]struct{}を利用した。
[]stringを集合として扱い比較したい場合には、それぞれの要素が過不足なく含まれていることを確認するためにforループをたくさん回すことになるだろう。

sample.go

package sample

import (
	"testing"

	"github.com/go-redis/redis/v7"
)

type RepositoryImpl struct {
	Client *redis.Client
}

func New(addr string) (*redis.Client, error) {
	client := redis.NewClient(&redis.Options{
		Addr: addr,
	})
	if err := client.Ping().Err(); err != nil {
		return nil, errors.Wrapf(err, "failed to ping redis server")
	}
	return client, nil
}

func (r RepositoryImpl) GetStringSet(key string) ([]string, error) {
	result, err := r.Client.SMembers(key).Result()
	if err != nil {
		return nil, err
	}
	return result, nil
}

sample_test.go

package sample

import (
	"testing"

	"github.com/go-redis/redis/v7"
)

func NewMockRedis(t *testing.T) *redis.Client {
	t.Helper()

	s, err := miniredis.Run()
	if err != nil {
		t.Fatalf("unexpected error while createing test redis server '%#v'", err)
	}

	client := redis.NewClient(&redis.Options{
		Addr: s.Addr(),
	})
	return client
}


func TestGetStringSet(t *testing.T) {
	client := NewMockRedis(t)
	r := RepositoryImpl{
		Client: client,
	}

	client.SAdd("testKey", "hoge", "fuga", "piyo")

	actual, err := r.GetStringSet("testKey")
	if err != nil {
		t.Fatalf("unexpected error while GetStringSet '%#v'", err)
	}
	expected, err := []string{"hoge", "fuga", "piyo"}

	// 順番が担保されないため、テストが失敗する
	if diff := cmp.Diff(actual, expected); diff != "" {
		t.Errorf("Diff: (-got +want)\n%s", diff)
	}

	// sliceではcmp.Diffで順序が考慮されてしまうのでSetに変換して比較する
	expectedSet := make(map[string]struct{})
	for _, v := range expected {
		expectedSet[v] = struct{}{}
	}
	actualSet := make(map[string]struct{})
	for _, v := range actual {
		actualSet[v] = struct{}{}
	}

	// 順番が担保されるため、テストが成功する
	if diff := cmp.Diff(actualSet, expectedSet); diff != "" {
		t.Errorf("Diff: (-got +want)\n%s", diff)
	}
}

参考

GoでSet型を実現する場合の選択肢