github.com/devfans/go-ethereum@v1.5.10-0.20170326212234-7419d0c38291/rpc/subscription_test.go (about)

     1  // Copyright 2016 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package rpc
    18  
    19  import (
    20  	"context"
    21  	"encoding/json"
    22  	"net"
    23  	"sync"
    24  	"testing"
    25  	"time"
    26  )
    27  
    28  type NotificationTestService struct {
    29  	mu           sync.Mutex
    30  	unsubscribed bool
    31  
    32  	gotHangSubscriptionReq  chan struct{}
    33  	unblockHangSubscription chan struct{}
    34  }
    35  
    36  func (s *NotificationTestService) Echo(i int) int {
    37  	return i
    38  }
    39  
    40  func (s *NotificationTestService) wasUnsubCallbackCalled() bool {
    41  	s.mu.Lock()
    42  	defer s.mu.Unlock()
    43  	return s.unsubscribed
    44  }
    45  
    46  func (s *NotificationTestService) Unsubscribe(subid string) {
    47  	s.mu.Lock()
    48  	s.unsubscribed = true
    49  	s.mu.Unlock()
    50  }
    51  
    52  func (s *NotificationTestService) SomeSubscription(ctx context.Context, n, val int) (*Subscription, error) {
    53  	notifier, supported := NotifierFromContext(ctx)
    54  	if !supported {
    55  		return nil, ErrNotificationsUnsupported
    56  	}
    57  
    58  	// by explicitly creating an subscription we make sure that the subscription id is send back to the client
    59  	// before the first subscription.Notify is called. Otherwise the events might be send before the response
    60  	// for the eth_subscribe method.
    61  	subscription := notifier.CreateSubscription()
    62  
    63  	go func() {
    64  		// test expects n events, if we begin sending event immediately some events
    65  		// will probably be dropped since the subscription ID might not be send to
    66  		// the client.
    67  		time.Sleep(5 * time.Second)
    68  		for i := 0; i < n; i++ {
    69  			if err := notifier.Notify(subscription.ID, val+i); err != nil {
    70  				return
    71  			}
    72  		}
    73  
    74  		select {
    75  		case <-notifier.Closed():
    76  			s.mu.Lock()
    77  			s.unsubscribed = true
    78  			s.mu.Unlock()
    79  		case <-subscription.Err():
    80  			s.mu.Lock()
    81  			s.unsubscribed = true
    82  			s.mu.Unlock()
    83  		}
    84  	}()
    85  
    86  	return subscription, nil
    87  }
    88  
    89  // HangSubscription blocks on s.unblockHangSubscription before
    90  // sending anything.
    91  func (s *NotificationTestService) HangSubscription(ctx context.Context, val int) (*Subscription, error) {
    92  	notifier, supported := NotifierFromContext(ctx)
    93  	if !supported {
    94  		return nil, ErrNotificationsUnsupported
    95  	}
    96  
    97  	s.gotHangSubscriptionReq <- struct{}{}
    98  	<-s.unblockHangSubscription
    99  	subscription := notifier.CreateSubscription()
   100  
   101  	go func() {
   102  		notifier.Notify(subscription.ID, val)
   103  	}()
   104  	return subscription, nil
   105  }
   106  
   107  func TestNotifications(t *testing.T) {
   108  	server := NewServer()
   109  	service := &NotificationTestService{}
   110  
   111  	if err := server.RegisterName("eth", service); err != nil {
   112  		t.Fatalf("unable to register test service %v", err)
   113  	}
   114  
   115  	clientConn, serverConn := net.Pipe()
   116  
   117  	go server.ServeCodec(NewJSONCodec(serverConn), OptionMethodInvocation|OptionSubscriptions)
   118  
   119  	out := json.NewEncoder(clientConn)
   120  	in := json.NewDecoder(clientConn)
   121  
   122  	n := 5
   123  	val := 12345
   124  	request := map[string]interface{}{
   125  		"id":      1,
   126  		"method":  "eth_subscribe",
   127  		"version": "2.0",
   128  		"params":  []interface{}{"someSubscription", n, val},
   129  	}
   130  
   131  	// create subscription
   132  	if err := out.Encode(request); err != nil {
   133  		t.Fatal(err)
   134  	}
   135  
   136  	var subid string
   137  	response := jsonSuccessResponse{Result: subid}
   138  	if err := in.Decode(&response); err != nil {
   139  		t.Fatal(err)
   140  	}
   141  
   142  	var ok bool
   143  	if _, ok = response.Result.(string); !ok {
   144  		t.Fatalf("expected subscription id, got %T", response.Result)
   145  	}
   146  
   147  	for i := 0; i < n; i++ {
   148  		var notification jsonNotification
   149  		if err := in.Decode(&notification); err != nil {
   150  			t.Fatalf("%v", err)
   151  		}
   152  
   153  		if int(notification.Params.Result.(float64)) != val+i {
   154  			t.Fatalf("expected %d, got %d", val+i, notification.Params.Result)
   155  		}
   156  	}
   157  
   158  	clientConn.Close() // causes notification unsubscribe callback to be called
   159  	time.Sleep(1 * time.Second)
   160  
   161  	if !service.wasUnsubCallbackCalled() {
   162  		t.Error("unsubscribe callback not called after closing connection")
   163  	}
   164  }