Add WSClient

pull/39/head
178inaba 2017-04-24 13:55:07 +09:00
parent ee6a6dbcc9
commit 010ec2eaf9
2 changed files with 30 additions and 21 deletions

View File

@ -9,6 +9,15 @@ import (
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
) )
// WSClient is a WebSocket client.
type WSClient struct {
websocket.Dialer
client *Client
}
// NewWSClient return WebSocket client.
func (c *Client) NewWSClient() *WSClient { return &WSClient{client: c} }
// Stream is a struct of data that flows in streaming. // Stream is a struct of data that flows in streaming.
type Stream struct { type Stream struct {
Event string `json:"event"` Event string `json:"event"`
@ -16,39 +25,39 @@ type Stream struct {
} }
// StreamingWSPublic return channel to read events on public using WebSocket. // StreamingWSPublic return channel to read events on public using WebSocket.
func (c *Client) StreamingWSPublic(ctx context.Context) (chan Event, error) { func (c *WSClient) StreamingWSPublic(ctx context.Context) (chan Event, error) {
return c.streamingWS(ctx, "public", "") return c.streamingWS(ctx, "public", "")
} }
// StreamingWSPublicLocal return channel to read events on public local using WebSocket. // StreamingWSPublicLocal return channel to read events on public local using WebSocket.
func (c *Client) StreamingWSPublicLocal(ctx context.Context) (chan Event, error) { func (c *WSClient) StreamingWSPublicLocal(ctx context.Context) (chan Event, error) {
return c.streamingWS(ctx, "public:local", "") return c.streamingWS(ctx, "public:local", "")
} }
// StreamingWSUser return channel to read events on home using WebSocket. // StreamingWSUser return channel to read events on home using WebSocket.
func (c *Client) StreamingWSUser(ctx context.Context) (chan Event, error) { func (c *WSClient) StreamingWSUser(ctx context.Context) (chan Event, error) {
return c.streamingWS(ctx, "user", "") return c.streamingWS(ctx, "user", "")
} }
// StreamingWSHashtag return channel to read events on tagged timeline using WebSocket. // StreamingWSHashtag return channel to read events on tagged timeline using WebSocket.
func (c *Client) StreamingWSHashtag(ctx context.Context, tag string) (chan Event, error) { func (c *WSClient) StreamingWSHashtag(ctx context.Context, tag string) (chan Event, error) {
return c.streamingWS(ctx, "hashtag", tag) return c.streamingWS(ctx, "hashtag", tag)
} }
// StreamingWSHashtagLocal return channel to read events on tagged local timeline using WebSocket. // StreamingWSHashtagLocal return channel to read events on tagged local timeline using WebSocket.
func (c *Client) StreamingWSHashtagLocal(ctx context.Context, tag string) (chan Event, error) { func (c *WSClient) StreamingWSHashtagLocal(ctx context.Context, tag string) (chan Event, error) {
return c.streamingWS(ctx, "hashtag:local", tag) return c.streamingWS(ctx, "hashtag:local", tag)
} }
func (c *Client) streamingWS(ctx context.Context, stream, tag string) (chan Event, error) { func (c *WSClient) streamingWS(ctx context.Context, stream, tag string) (chan Event, error) {
params := url.Values{} params := url.Values{}
params.Set("access_token", c.config.AccessToken) params.Set("access_token", c.client.config.AccessToken)
params.Set("stream", stream) params.Set("stream", stream)
if tag != "" { if tag != "" {
params.Set("tag", tag) params.Set("tag", tag)
} }
u, err := changeWebSocketScheme(c.config.Server) u, err := changeWebSocketScheme(c.client.config.Server)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -68,7 +77,7 @@ func (c *Client) streamingWS(ctx context.Context, stream, tag string) (chan Even
return q, nil return q, nil
} }
func (c *Client) handleWS(ctx context.Context, rawurl string, q chan Event) error { func (c *WSClient) handleWS(ctx context.Context, rawurl string, q chan Event) error {
conn, err := c.dialRedirect(rawurl) conn, err := c.dialRedirect(rawurl)
if err != nil { if err != nil {
q <- &ErrorEvent{err: err} q <- &ErrorEvent{err: err}
@ -122,7 +131,7 @@ func (c *Client) handleWS(ctx context.Context, rawurl string, q chan Event) erro
return nil return nil
} }
func (c *Client) dialRedirect(rawurl string) (conn *websocket.Conn, err error) { func (c *WSClient) dialRedirect(rawurl string) (conn *websocket.Conn, err error) {
for { for {
conn, rawurl, err = c.dial(rawurl) conn, rawurl, err = c.dial(rawurl)
if err != nil { if err != nil {
@ -133,7 +142,7 @@ func (c *Client) dialRedirect(rawurl string) (conn *websocket.Conn, err error) {
} }
} }
func (c *Client) dial(rawurl string) (*websocket.Conn, string, error) { func (c *WSClient) dial(rawurl string) (*websocket.Conn, string, error) {
conn, resp, err := c.Dial(rawurl, nil) conn, resp, err := c.Dial(rawurl, nil)
if err != nil && err != websocket.ErrBadHandshake { if err != nil && err != websocket.ErrBadHandshake {
return nil, "", err return nil, "", err

View File

@ -14,7 +14,7 @@ func TestStreamingWSPublic(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(wsMock)) ts := httptest.NewServer(http.HandlerFunc(wsMock))
defer ts.Close() defer ts.Close()
client := NewClient(&Config{Server: ts.URL}) client := NewClient(&Config{Server: ts.URL}).NewWSClient()
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
q, err := client.StreamingWSPublic(ctx) q, err := client.StreamingWSPublic(ctx)
if err != nil { if err != nil {
@ -28,7 +28,7 @@ func TestStreamingWSPublicLocal(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(wsMock)) ts := httptest.NewServer(http.HandlerFunc(wsMock))
defer ts.Close() defer ts.Close()
client := NewClient(&Config{Server: ts.URL}) client := NewClient(&Config{Server: ts.URL}).NewWSClient()
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
q, err := client.StreamingWSPublicLocal(ctx) q, err := client.StreamingWSPublicLocal(ctx)
if err != nil { if err != nil {
@ -42,7 +42,7 @@ func TestStreamingWSUser(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(wsMock)) ts := httptest.NewServer(http.HandlerFunc(wsMock))
defer ts.Close() defer ts.Close()
client := NewClient(&Config{Server: ts.URL}) client := NewClient(&Config{Server: ts.URL}).NewWSClient()
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
q, err := client.StreamingWSUser(ctx) q, err := client.StreamingWSUser(ctx)
if err != nil { if err != nil {
@ -56,7 +56,7 @@ func TestStreamingWSHashtag(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(wsMock)) ts := httptest.NewServer(http.HandlerFunc(wsMock))
defer ts.Close() defer ts.Close()
client := NewClient(&Config{Server: ts.URL}) client := NewClient(&Config{Server: ts.URL}).NewWSClient()
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
q, err := client.StreamingWSHashtag(ctx, "zzz") q, err := client.StreamingWSHashtag(ctx, "zzz")
if err != nil { if err != nil {
@ -70,7 +70,7 @@ func TestStreamingWSHashtagLocal(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(wsMock)) ts := httptest.NewServer(http.HandlerFunc(wsMock))
defer ts.Close() defer ts.Close()
client := NewClient(&Config{Server: ts.URL}) client := NewClient(&Config{Server: ts.URL}).NewWSClient()
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
q, err := client.StreamingWSHashtagLocal(ctx, "zzz") q, err := client.StreamingWSHashtagLocal(ctx, "zzz")
if err != nil { if err != nil {
@ -155,13 +155,13 @@ func TestStreamingWS(t *testing.T) {
ts := httptest.NewServer(http.HandlerFunc(wsMock)) ts := httptest.NewServer(http.HandlerFunc(wsMock))
defer ts.Close() defer ts.Close()
client := NewClient(&Config{Server: ":"}) client := NewClient(&Config{Server: ":"}).NewWSClient()
_, err := client.StreamingWSPublicLocal(context.Background()) _, err := client.StreamingWSPublicLocal(context.Background())
if err == nil { if err == nil {
t.Fatalf("should be fail: %v", err) t.Fatalf("should be fail: %v", err)
} }
client = NewClient(&Config{Server: ts.URL}) client = NewClient(&Config{Server: ts.URL}).NewWSClient()
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
cancel() cancel()
q, err := client.StreamingWSPublicLocal(ctx) q, err := client.StreamingWSPublicLocal(ctx)
@ -198,7 +198,7 @@ func TestHandleWS(t *testing.T) {
defer ts.Close() defer ts.Close()
q := make(chan Event) q := make(chan Event)
client := NewClient(&Config{}) client := NewClient(&Config{}).NewWSClient()
go func() { go func() {
e := <-q e := <-q
@ -234,7 +234,7 @@ func TestHandleWS(t *testing.T) {
} }
func TestDialRedirect(t *testing.T) { func TestDialRedirect(t *testing.T) {
client := NewClient(&Config{}) client := NewClient(&Config{}).NewWSClient()
_, err := client.dialRedirect(":") _, err := client.dialRedirect(":")
if err == nil { if err == nil {
t.Fatalf("should be fail: %v", err) t.Fatalf("should be fail: %v", err)
@ -254,7 +254,7 @@ func TestDial(t *testing.T) {
})) }))
defer ts.Close() defer ts.Close()
client := NewClient(&Config{}) client := NewClient(&Config{}).NewWSClient()
_, _, err := client.dial(":") _, _, err := client.dial(":")
if err == nil { if err == nil {
t.Fatalf("should be fail: %v", err) t.Fatalf("should be fail: %v", err)