package test | |
import ( | |
"net" | |
"github.com/coredns/coredns/plugin/test" | |
"github.com/miekg/dns" | |
) | |
// ResponseWriter extends test.ResponseWriter and record useful calls to methods. | |
type ResponseWriter struct { | |
test.ResponseWriter | |
writeMsgCallCount uint64 | |
} | |
// ResponseWriterCustomRemote is a ResponseWriter which can receive a custom | |
// remote address. | |
type ResponseWriterCustomRemote struct { | |
ResponseWriter | |
RemoteIP string | |
} | |
// WriteMsg implement dns.ResponseWriter interface. | |
func (t *ResponseWriter) WriteMsg(m *dns.Msg) error { t.writeMsgCallCount++; return nil } | |
// GetWriteMsgCallCount returns the number of WriteMsg calls that were made. | |
func (t *ResponseWriter) GetWriteMsgCallCount() uint64 { return t.writeMsgCallCount } | |
// RemoteAddr returns a net.Addr of the remote address connecting to the server. | |
func (t *ResponseWriterCustomRemote) RemoteAddr() net.Addr { | |
return &net.UDPAddr{IP: net.ParseIP(t.RemoteIP), Port: 40212, Zone: ""} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment