pgtype/enum_type_test.go at master · jackc/pgtype · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
package pgtype_test

import (
	"bytes"
	"context"
	"testing"

	"github.com/jackc/pgtype"
	"github.com/jackc/pgtype/testutil"
	"github.com/jackc/pgx/v4"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func setupEnum(t *testing.T, conn *pgx.Conn) *pgtype.EnumType {
	_, err := conn.Exec(context.Background(), "drop type if exists pgtype_enum_color;")
	require.NoError(t, err)

	_, err = conn.Exec(context.Background(), "create type pgtype_enum_color as enum ('blue', 'green', 'purple');")
	require.NoError(t, err)

	var oid uint32
	err = conn.QueryRow(context.Background(), "select oid from pg_type where typname=$1;", "pgtype_enum_color").Scan(&oid)
	require.NoError(t, err)

	et := pgtype.NewEnumType("pgtype_enum_color", []string{"blue", "green", "purple"})
	conn.ConnInfo().RegisterDataType(pgtype.DataType{Value: et, Name: "pgtype_enum_color", OID: oid})

	return et
}

func cleanupEnum(t *testing.T, conn *pgx.Conn) {
	_, err := conn.Exec(context.Background(), "drop type if exists pgtype_enum_color;")
	require.NoError(t, err)
}

func TestEnumTypeTranscode(t *testing.T) {
	conn := testutil.MustConnectPgx(t)
	defer testutil.MustCloseContext(t, conn)

	setupEnum(t, conn)
	defer cleanupEnum(t, conn)

	var dst string
	err := conn.QueryRow(context.Background(), "select $1::pgtype_enum_color", "blue").Scan(&dst)
	require.NoError(t, err)
	require.EqualValues(t, "blue", dst)
}

func TestEnumTypeSet(t *testing.T) {
	conn := testutil.MustConnectPgx(t)
	defer testutil.MustCloseContext(t, conn)

	enumType := setupEnum(t, conn)
	defer cleanupEnum(t, conn)

	successfulTests := []struct {
		source interface{}
		result interface{}
	}{
		{source: "blue", result: "blue"},
		{source: _string("green"), result: "green"},
		{source: (*string)(nil), result: nil},
	}

	for i, tt := range successfulTests {
		err := enumType.Set(tt.source)
		assert.NoErrorf(t, err, "%d", i)
		assert.Equalf(t, tt.result, enumType.Get(), "%d", i)
	}
}

func TestEnumTypeAssignTo(t *testing.T) {
	conn := testutil.MustConnectPgx(t)
	defer testutil.MustCloseContext(t, conn)

	enumType := setupEnum(t, conn)
	defer cleanupEnum(t, conn)

	{
		var s string

		err := enumType.Set("blue")
		require.NoError(t, err)

		err = enumType.AssignTo(&s)
		require.NoError(t, err)

		assert.EqualValues(t, "blue", s)
	}

	{
		var ps *string

		err := enumType.Set("blue")
		require.NoError(t, err)

		err = enumType.AssignTo(&ps)
		require.NoError(t, err)

		assert.EqualValues(t, "blue", *ps)
	}

	{
		var ps *string

		err := enumType.Set(nil)
		require.NoError(t, err)

		err = enumType.AssignTo(&ps)
		require.NoError(t, err)

		assert.EqualValues(t, (*string)(nil), ps)
	}

	var buf []byte
	bytesTests := []struct {
		src      interface{}
		dst      *[]byte
		expected []byte
	}{
		{src: "blue", dst: &buf, expected: []byte("blue")},
		{src: nil, dst: &buf, expected: nil},
	}

	for i, tt := range bytesTests {
		err := enumType.Set(tt.src)
		require.NoError(t, err, "%d", i)

		err = enumType.AssignTo(tt.dst)
		require.NoError(t, err, "%d", i)

		if !bytes.Equal(*tt.dst, tt.expected) {
			t.Errorf("%d: expected %v to assign %v, but result was %v", i, tt.src, tt.expected, tt.dst)
		}
	}

	{
		var s string

		err := enumType.Set(nil)
		require.NoError(t, err)

		err = enumType.AssignTo(&s)
		require.Error(t, err)
	}

}