-
Notifications
You must be signed in to change notification settings - Fork 42
/
mock.go
91 lines (77 loc) · 2.72 KB
/
mock.go
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
package mock
import (
"context"
"io"
"net"
"github.com/cyberark/secretless-broker/internal/plugin/connectors/tcp/mssql/types"
mssql "github.com/denisenkom/go-mssqldb"
)
// NewSuccessfulMSSQLConnectorCtor returns an MSSQLConnectorCtor that always
// succeeds.
func NewSuccessfulMSSQLConnectorCtor(
fn types.MSSQLConnectorFunc,
) types.MSSQLConnectorCtor {
return func(dsn string) (types.MSSQLConnector, error) {
return types.MSSQLConnector(fn), nil
}
}
// NewFailingMSSQLConnectorCtor returns an MSSQLConnectorCtor that always
// returns the specified error.
func NewFailingMSSQLConnectorCtor(err error) types.MSSQLConnectorCtor {
return func(dsn string) (types.MSSQLConnector, error) {
return nil, err
}
}
// NewSuccessfulMSSQLConnector returns an MSSQLConnector double whose Connect
// method always succeeds.
func NewSuccessfulMSSQLConnector(
fn func(context.Context) (net.Conn, error),
) types.MSSQLConnector {
return types.MSSQLConnectorFunc(fn)
}
// NewFailingMSSQLConnector returns an MSSQLConnector double whose Connect
// method always fails.
func NewFailingMSSQLConnector(err error) types.MSSQLConnector {
rawFunc := func(context.Context) (net.Conn, error) {
return nil, err
}
return types.MSSQLConnectorFunc(rawFunc)
}
// SuccessfulReadPreloginRequest is a double for a ReadPreloginRequestFunc that always
// succeeds.
func SuccessfulReadPreloginRequest(io.ReadWriteCloser) (map[uint8][]byte, error) {
return nil, nil
}
// SuccessfulWritePreloginResponse is a double for a WritePreloginResponseFunc that always
// succeeds.
func SuccessfulWritePreloginResponse(io.ReadWriteCloser, map[uint8][]byte) error {
return nil
}
// SuccessfulReadLoginRequest is a double for a ReadLoginRequestFunc that always succeeds.
func SuccessfulReadLoginRequest(io.ReadWriteCloser) (*mssql.LoginRequest, error) {
return &mssql.LoginRequest{}, nil
}
// SuccessfulWriteError is a double for a WriteErrorFunc that always succeeds.
func SuccessfulWriteError(io.ReadWriteCloser, mssql.Error) error {
return nil
}
// NewNetConn returns a net.Conn double whose behavior we can control.
func NewNetConn(errOnWrite error) *NetConn {
return &NetConn{errOnWrite: errOnWrite}
}
// NetConn acts as a double of a true network connection, ie, a net.Conn.
// TODO: This will need to be upgraded to have more granularity. For example,
// to handle cases where sending the authentication OK message works, but
// sending an error fails. Etc.
type NetConn struct {
net.Conn
errOnWrite error
}
// Write "writes" bytes to our fake net.Conn.
func (n *NetConn) Write([]byte) (numBytes int, err error) {
return 1, n.errOnWrite
}
// FakeTdsBufferCtor returns the ReadWriteCloser passed in.
func FakeTdsBufferCtor(r io.ReadWriteCloser) io.ReadWriteCloser {
return r
}