forked from aws/aws-sdk-go
-
Notifications
You must be signed in to change notification settings - Fork 3
/
testing.go
101 lines (85 loc) · 2.77 KB
/
testing.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
92
93
94
95
96
97
98
99
100
101
package eventstreamtest
import (
"bytes"
"fmt"
"io"
"net/http"
"net/http/httptest"
"testing"
"github.com/journeymidnight/aws-sdk-go/aws"
"github.com/journeymidnight/aws-sdk-go/aws/session"
"github.com/journeymidnight/aws-sdk-go/awstesting/unit"
"github.com/journeymidnight/aws-sdk-go/private/protocol"
"github.com/journeymidnight/aws-sdk-go/private/protocol/eventstream"
"github.com/journeymidnight/aws-sdk-go/private/protocol/eventstream/eventstreamapi"
)
// ServeEventStream provides serving EventStream messages from a HTTP server to
// the client. The events are sent sequentially to the client without delay.
type ServeEventStream struct {
T *testing.T
Events []eventstream.Message
}
func (s ServeEventStream) ServeHTTP(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
encoder := eventstream.NewEncoder(flushWriter{w})
for _, event := range s.Events {
encoder.Encode(event)
}
}
// SetupEventStreamSession creates a HTTP server SDK session for communicating
// with that server to be used for EventStream APIs. If HTTP/2 is enabled the
// server/client will only attempt to use HTTP/2.
func SetupEventStreamSession(
t *testing.T, handler http.Handler, h2 bool,
) (sess *session.Session, cleanupFn func(), err error) {
server := httptest.NewUnstartedServer(handler)
client := setupServer(server, h2)
cleanupFn = func() {
server.Close()
}
sess, err = session.NewSession(unit.Session.Config, &aws.Config{
Endpoint: &server.URL,
DisableParamValidation: aws.Bool(true),
HTTPClient: client,
// LogLevel: aws.LogLevel(aws.LogDebugWithEventStreamBody),
})
if err != nil {
return nil, nil, err
}
return sess, cleanupFn, nil
}
type flushWriter struct {
w io.Writer
}
func (fw flushWriter) Write(p []byte) (n int, err error) {
n, err = fw.w.Write(p)
if f, ok := fw.w.(http.Flusher); ok {
f.Flush()
}
return
}
// MarshalEventPayload marshals a SDK API shape into its associated wire
// protocol payload.
func MarshalEventPayload(
payloadMarshaler protocol.PayloadMarshaler,
v interface{},
) []byte {
var w bytes.Buffer
err := payloadMarshaler.MarshalPayload(&w, v)
if err != nil {
panic(fmt.Sprintf("failed to marshal event %T, %v", v, v))
}
return w.Bytes()
}
// EventMessageTypeHeader is an event message type header for specifying an
// event is an message type.
var EventMessageTypeHeader = eventstream.Header{
Name: eventstreamapi.MessageTypeHeader,
Value: eventstream.StringValue(eventstreamapi.EventMessageType),
}
// EventExceptionTypeHeader is an event exception type header for specifying an
// event is an exeption type.
var EventExceptionTypeHeader = eventstream.Header{
Name: eventstreamapi.MessageTypeHeader,
Value: eventstream.StringValue(eventstreamapi.ExceptionMessageType),
}