pgsnap/custom_expectation.go at main · egon12/pgsnap · 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
package pgsnap

import (
	"fmt"
	"reflect"

	"github.com/jackc/pgproto3/v2"
)

type (
	// expectParseMessage is a custom expectation for pgx that ignore Name
	expectParseMessage struct{ want *pgproto3.Parse }

	// expectDescribeMessage is a custom expectation for pgx that ignore Name
	expectDescribeMessage struct{ want *pgproto3.Describe }

	// expectBindMessage is a custom expectation for pgx that ignore PreparedStatement
	expectBindMessage struct{ want *pgproto3.Bind }
)

func (e *expectParseMessage) Step(backend *pgproto3.Backend) error {
	msg, err := backend.Receive()
	if err != nil {
		return err
	}
	return e.compare(msg)
}

// we ignore and m.Name, because it's inconsisten in pgx
func (e *expectParseMessage) compare(msg pgproto3.FrontendMessage) error {
	m, ok := msg.(*pgproto3.Parse)
	if !ok {
		return fmt.Errorf("msg => %T, want => %T", msg, e.want)
	}

	if m.Query != e.want.Query {
		return fmt.Errorf("msg => query: %s, want => query: %s", m.Query, e.want.Query)
	}

	if !reflect.DeepEqual(m.ParameterOIDs, e.want.ParameterOIDs) {
		return fmt.Errorf("msg => ParameterOIDs: %v, want => ParameterOIDs: %v", m.ParameterOIDs, e.want.ParameterOIDs)
	}

	return nil
}

func (e *expectDescribeMessage) Step(backend *pgproto3.Backend) error {
	msg, err := backend.Receive()
	if err != nil {
		return err
	}
	return e.compare(msg)
}

// we ignore and m.Name, because it's inconsisten in pgx
func (e *expectDescribeMessage) compare(msg pgproto3.FrontendMessage) error {
	m, ok := msg.(*pgproto3.Describe)
	if !ok {
		return fmt.Errorf("msg => %T, want => %T", msg, e.want)
	}

	if m.ObjectType != e.want.ObjectType {
		return fmt.Errorf("msg => ObjectType: %s, want => ObjectType: %s", string(m.ObjectType), string(e.want.ObjectType))
	}

	return nil
}

func (e *expectBindMessage) Step(backend *pgproto3.Backend) error {
	msg, err := backend.Receive()
	if err != nil {
		return err
	}

	return e.compare(msg)
}

func (e *expectBindMessage) compare(msg pgproto3.FrontendMessage) error {
	m, ok := msg.(*pgproto3.Bind)
	if !ok {
		return fmt.Errorf("msg => %T, want => %T", msg, e.want)
	}

	// we ignore and m.Name, because it's inconsisten in pgx
	if m.DestinationPortal != e.want.DestinationPortal {
		return fmt.Errorf(
			"msg => DestinationPortal: %s, want => DestinationPortal: %s",
			m.DestinationPortal,
			e.want.DestinationPortal,
		)
	}

	if !reflect.DeepEqual(m.ParameterFormatCodes, e.want.ParameterFormatCodes) {
		return fmt.Errorf(
			"msg => ParameterFormatCodes: %v, want => ParameterFormatCodes: %v",
			m.ParameterFormatCodes,
			e.want.ParameterFormatCodes,
		)
	}

	if !reflect.DeepEqual(m.ResultFormatCodes, e.want.ResultFormatCodes) {
		return fmt.Errorf(
			"msg => ResultFormatCodes: %v, want => ResultFormatCodes: %v",
			m.ResultFormatCodes,
			e.want.ResultFormatCodes,
		)
	}

	if !reflect.DeepEqual(m.Parameters, e.want.Parameters) {
		if len(m.Parameters) == 0 && len(e.want.Parameters) == 0 {
			return nil
		}

		return fmt.Errorf(
			"msg => Parameters: %v, want => Parameters: %v",
			m.Parameters,
			e.want.Parameters,
		)
	}

	return nil
}