From 38d18ca11a5647889d1ad8ca1b3d53ed9627f6e1 Mon Sep 17 00:00:00 2001
From: Alexander Simmerl <a.simmerl@gmail.com>
Date: Tue, 6 Feb 2018 17:39:10 +0100
Subject: [PATCH] Harden tests

---
 types/priv_validator/socket.go      |  6 +--
 types/priv_validator/socket_test.go | 77 +++++++++++++++++++++++++++--
 2 files changed, 75 insertions(+), 8 deletions(-)

diff --git a/types/priv_validator/socket.go b/types/priv_validator/socket.go
index 26030346..1fd45ed6 100644
--- a/types/priv_validator/socket.go
+++ b/types/priv_validator/socket.go
@@ -137,7 +137,7 @@ func (pvsc *PrivValidatorSocketClient) SignVote(chainID string, vote *types.Vote
 		return err
 	}
 
-	*vote = *res.(SignVoteMsg).Vote
+	*vote = *res.(*SignVoteMsg).Vote
 
 	return nil
 }
@@ -154,7 +154,7 @@ func (pvsc *PrivValidatorSocketClient) SignProposal(chainID string, proposal *ty
 		return err
 	}
 
-	*proposal = *res.(SignProposalMsg).Proposal
+	*proposal = *res.(*SignProposalMsg).Proposal
 
 	return nil
 }
@@ -171,7 +171,7 @@ func (pvsc *PrivValidatorSocketClient) SignHeartbeat(chainID string, heartbeat *
 		return err
 	}
 
-	*heartbeat = *res.(SignHeartbeatMsg).Heartbeat
+	*heartbeat = *res.(*SignHeartbeatMsg).Heartbeat
 
 	return nil
 }
diff --git a/types/priv_validator/socket_test.go b/types/priv_validator/socket_test.go
index 88e19009..725d5af6 100644
--- a/types/priv_validator/socket_test.go
+++ b/types/priv_validator/socket_test.go
@@ -1,13 +1,14 @@
 package types
 
 import (
-	"reflect"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	crypto "github.com/tendermint/go-crypto"
+	data "github.com/tendermint/go-wire/data"
 	"github.com/tendermint/tendermint/types"
 	"github.com/tendermint/tmlibs/log"
 )
@@ -15,7 +16,7 @@ import (
 func TestPrivValidatorSocketServer(t *testing.T) {
 	var (
 		assert, require = assert.New(t), require.New(t)
-		chainID         = "test-chain"
+		chainID         = "test-chain-secret"
 		logger          = log.TestingLogger()
 		signer          = types.GenSigner()
 		clientPrivKey   = crypto.GenPrivKeyEd25519()
@@ -49,7 +50,73 @@ func TestPrivValidatorSocketServer(t *testing.T) {
 
 	assert.True(pvsc.IsRunning())
 
-	if have, want := pvsc.PubKey(), pvss.privVal.PubKey(); !reflect.DeepEqual(have, want) {
-		t.Errorf("have %v, want %v", have, want)
-	}
+	assert.Equal(pvsc.Address(), data.Bytes(pvss.privVal.PubKey().Address()))
+	assert.Equal(pvsc.PubKey(), pvss.privVal.PubKey())
+
+	err = pvsc.SignProposal(chainID, &types.Proposal{
+		Timestamp: time.Now(),
+	})
+	require.Nil(err)
+
+	err = pvsc.SignVote(chainID, &types.Vote{
+		Timestamp: time.Now(),
+		Type:      types.VoteTypePrecommit,
+	})
+	require.Nil(err)
+
+	err = pvsc.SignHeartbeat(chainID, &types.Heartbeat{})
+	require.Nil(err)
+}
+
+func TestPrivValidatorSocketServerWithoutSecret(t *testing.T) {
+	var (
+		assert, require = assert.New(t), require.New(t)
+		chainID         = "test-chain-secret"
+		logger          = log.TestingLogger()
+		signer          = types.GenSigner()
+		privVal         = NewTestPrivValidator(signer)
+		pvss            = NewPrivValidatorSocketServer(
+			logger,
+			chainID,
+			"127.0.0.1:0",
+			1,
+			privVal,
+			nil,
+		)
+	)
+
+	err := pvss.Start()
+	require.Nil(err)
+	defer pvss.Stop()
+
+	assert.True(pvss.IsRunning())
+
+	pvsc := NewPrivValidatorSocketClient(
+		logger,
+		pvss.listener.Addr().String(),
+		nil,
+	)
+
+	err = pvsc.Start()
+	require.Nil(err)
+	defer pvsc.Stop()
+
+	assert.True(pvsc.IsRunning())
+
+	assert.Equal(pvsc.Address(), data.Bytes(pvss.privVal.PubKey().Address()))
+	assert.Equal(pvsc.PubKey(), pvss.privVal.PubKey())
+
+	err = pvsc.SignProposal(chainID, &types.Proposal{
+		Timestamp: time.Now(),
+	})
+	require.Nil(err)
+
+	err = pvsc.SignVote(chainID, &types.Vote{
+		Timestamp: time.Now(),
+		Type:      types.VoteTypePrecommit,
+	})
+	require.Nil(err)
+
+	err = pvsc.SignHeartbeat(chainID, &types.Heartbeat{})
+	require.Nil(err)
 }