infra/proxyd/integration_tests/util_test.go
2024-06-11 14:17:04 -05:00

192 lines
3.9 KiB
Go

package integration_tests
import (
"bytes"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"testing"
"time"
"github.com/BurntSushi/toml"
"github.com/gorilla/websocket"
"github.com/stretchr/testify/require"
"golang.org/x/exp/slog"
"github.com/ethereum/go-ethereum/log"
"github.com/ethereum-optimism/optimism/proxyd"
)
type ProxydHTTPClient struct {
url string
headers http.Header
}
func NewProxydClient(url string) *ProxydHTTPClient {
return NewProxydClientWithHeaders(url, make(http.Header))
}
func NewProxydClientWithHeaders(url string, headers http.Header) *ProxydHTTPClient {
clonedHeaders := headers.Clone()
clonedHeaders.Set("Content-Type", "application/json")
return &ProxydHTTPClient{
url: url,
headers: clonedHeaders,
}
}
func (p *ProxydHTTPClient) SendRPC(method string, params []interface{}) ([]byte, int, error) {
rpcReq := NewRPCReq("999", method, params)
body, err := json.Marshal(rpcReq)
if err != nil {
panic(err)
}
return p.SendRequest(body)
}
func (p *ProxydHTTPClient) SendBatchRPC(reqs ...*proxyd.RPCReq) ([]byte, int, error) {
body, err := json.Marshal(reqs)
if err != nil {
panic(err)
}
return p.SendRequest(body)
}
func (p *ProxydHTTPClient) SendRequest(body []byte) ([]byte, int, error) {
req, err := http.NewRequest("POST", p.url, bytes.NewReader(body))
if err != nil {
panic(err)
}
req.Header = p.headers
res, err := http.DefaultClient.Do(req)
if err != nil {
return nil, -1, err
}
defer res.Body.Close()
code := res.StatusCode
resBody, err := io.ReadAll(res.Body)
if err != nil {
panic(err)
}
return resBody, code, nil
}
func RequireEqualJSON(t *testing.T, expected []byte, actual []byte) {
expJSON := canonicalizeJSON(t, expected)
actJSON := canonicalizeJSON(t, actual)
require.Equal(t, string(expJSON), string(actJSON))
}
func canonicalizeJSON(t *testing.T, in []byte) []byte {
var any interface{}
if in[0] == '[' {
any = make([]interface{}, 0)
} else {
any = make(map[string]interface{})
}
err := json.Unmarshal(in, &any)
require.NoError(t, err)
out, err := json.Marshal(any)
require.NoError(t, err)
return out
}
func ReadConfig(name string) *proxyd.Config {
config := new(proxyd.Config)
_, err := toml.DecodeFile(fmt.Sprintf("testdata/%s.toml", name), config)
if err != nil {
panic(err)
}
return config
}
func NewRPCReq(id string, method string, params []interface{}) *proxyd.RPCReq {
jsonParams, err := json.Marshal(params)
if err != nil {
panic(err)
}
return &proxyd.RPCReq{
JSONRPC: proxyd.JSONRPCVersion,
Method: method,
Params: jsonParams,
ID: []byte(id),
}
}
type ProxydWSClient struct {
conn *websocket.Conn
msgCB ProxydWSClientOnMessage
closeCB ProxydWSClientOnClose
}
type WSMessage struct {
Type int
Body []byte
}
type (
ProxydWSClientOnMessage func(msgType int, data []byte)
ProxydWSClientOnClose func(err error)
)
func NewProxydWSClient(
url string,
msgCB ProxydWSClientOnMessage,
closeCB ProxydWSClientOnClose,
) (*ProxydWSClient, error) {
conn, _, err := websocket.DefaultDialer.Dial(url, nil) // nolint:bodyclose
if err != nil {
return nil, err
}
c := &ProxydWSClient{
conn: conn,
msgCB: msgCB,
closeCB: closeCB,
}
go c.readPump()
return c, nil
}
func (h *ProxydWSClient) readPump() {
for {
mType, msg, err := h.conn.ReadMessage()
if err != nil {
if h.closeCB != nil {
h.closeCB(err)
}
return
}
if h.msgCB != nil {
h.msgCB(mType, msg)
}
}
}
func (h *ProxydWSClient) HardClose() {
h.conn.Close()
}
func (h *ProxydWSClient) SoftClose() error {
return h.WriteMessage(websocket.CloseMessage, nil)
}
func (h *ProxydWSClient) WriteMessage(msgType int, msg []byte) error {
return h.conn.WriteMessage(msgType, msg)
}
func (h *ProxydWSClient) WriteControlMessage(msgType int, msg []byte) error {
return h.conn.WriteControl(msgType, msg, time.Now().Add(time.Minute))
}
func InitLogger() {
log.SetDefault(log.NewLogger(slog.NewJSONHandler(
os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug})))
}