From bed044b5a1a0ee89a492496e01a303f4c4f532e9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Magiera?= <magik6k@gmail.com>
Date: Tue, 2 Jul 2019 15:20:33 +0200
Subject: [PATCH] Fix TestCtx

---
 rpclib/rpcServer.go |  2 +-
 rpclib/rpc_test.go  | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/rpclib/rpcServer.go b/rpclib/rpcServer.go
index fde02bd1e..967ca1edf 100644
--- a/rpclib/rpcServer.go
+++ b/rpclib/rpcServer.go
@@ -77,7 +77,7 @@ func (s *RPCServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	handler, ok := s.methods[req.Method]
 	if !ok {
-		fmt.Println("rpcserver: unknown method")
+		fmt.Printf("rpcserver: unknown method %s\n", req.Method)
 		w.WriteHeader(500)
 		return
 	}
diff --git a/rpclib/rpc_test.go b/rpclib/rpc_test.go
index 84686f2cf..df628c401 100644
--- a/rpclib/rpc_test.go
+++ b/rpclib/rpc_test.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"net/http/httptest"
 	"strconv"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -167,12 +168,12 @@ func TestRPC(t *testing.T) {
 
 type CtxHandler struct {
 	cancelled bool
-	i         int
+	i         int32
 }
 
 func (h *CtxHandler) Test(ctx context.Context) {
 	timeout := time.After(300 * time.Millisecond)
-	h.i++
+	atomic.AddInt32(&h.i, 1)
 
 	select {
 	case <-timeout:
@@ -182,13 +183,12 @@ func (h *CtxHandler) Test(ctx context.Context) {
 }
 
 func TestCtx(t *testing.T) {
-	t.SkipNow()
 	// setup server
 
 	serverHandler := &CtxHandler{}
 
 	rpcServer := NewServer()
-	rpcServer.Register("SimpleServerHandler", serverHandler)
+	rpcServer.Register("CtxHandler", serverHandler)
 
 	// httptest stuff
 	testServ := httptest.NewServer(rpcServer)
@@ -201,7 +201,7 @@ func TestCtx(t *testing.T) {
 	}
 	NewClient(testServ.URL, "CtxHandler", &client)
 
-	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 	defer cancel()
 
 	client.Test(ctx)
@@ -217,7 +217,7 @@ func TestCtx(t *testing.T) {
 	NewClient(testServ.URL, "CtxHandler", &noCtxClient)
 
 	noCtxClient.Test()
-	if serverHandler.cancelled || serverHandler.i != 2 {
+	if serverHandler.cancelled || atomic.LoadInt32(&serverHandler.i) != 2 {
 		t.Error("wrong serverHandler state")
 	}
 }