From 6aa85b642e4691dc9c57b2ba2d23a46972378175 Mon Sep 17 00:00:00 2001
From: Ethan Buchman <ethan@coinculture.info>
Date: Thu, 3 Nov 2016 19:50:57 -0400
Subject: [PATCH] BeginBlock

---
 client/socket_client.go                 |  2 +
 example/chain_aware/chain_aware_app.go  | 75 +++++++++++++++++++++++++
 example/chain_aware/chain_aware_test.go | 50 +++++++++++++++++
 example/dummy/dummy.go                  |  1 +
 server/socket_server.go                 |  7 +++
 5 files changed, 135 insertions(+)
 create mode 100644 example/chain_aware/chain_aware_app.go
 create mode 100644 example/chain_aware/chain_aware_test.go

diff --git a/client/socket_client.go b/client/socket_client.go
index 0b7be91d9..f650bd5c9 100644
--- a/client/socket_client.go
+++ b/client/socket_client.go
@@ -432,6 +432,8 @@ func resMatchesReq(req *types.Request, res *types.Response) (ok bool) {
 		_, ok = res.Value.(*types.Response_Query)
 	case *types.Request_InitChain:
 		_, ok = res.Value.(*types.Response_InitChain)
+	case *types.Request_BeginBlock:
+		_, ok = res.Value.(*types.Response_BeginBlock)
 	case *types.Request_EndBlock:
 		_, ok = res.Value.(*types.Response_EndBlock)
 	}
diff --git a/example/chain_aware/chain_aware_app.go b/example/chain_aware/chain_aware_app.go
new file mode 100644
index 000000000..df18891b5
--- /dev/null
+++ b/example/chain_aware/chain_aware_app.go
@@ -0,0 +1,75 @@
+package main
+
+import (
+	"flag"
+
+	. "github.com/tendermint/go-common"
+	"github.com/tendermint/tmsp/server"
+	"github.com/tendermint/tmsp/types"
+)
+
+func main() {
+
+	addrPtr := flag.String("addr", "tcp://0.0.0.0:46658", "Listen address")
+	tmspPtr := flag.String("tmsp", "socket", "socket | grpc")
+	flag.Parse()
+
+	// Start the listener
+	_, err := server.NewServer(*addrPtr, *tmspPtr, NewChainAwareApplication())
+	if err != nil {
+		Exit(err.Error())
+	}
+
+	// Wait forever
+	TrapSignal(func() {
+		// Cleanup
+	})
+
+}
+
+type ChainAwareApplication struct {
+	beginCount int
+	endCount   int
+}
+
+func NewChainAwareApplication() *ChainAwareApplication {
+	return &ChainAwareApplication{}
+}
+
+func (app *ChainAwareApplication) Info() string {
+	return "nil"
+}
+
+func (app *ChainAwareApplication) SetOption(key string, value string) (log string) {
+	return ""
+}
+
+func (app *ChainAwareApplication) AppendTx(tx []byte) types.Result {
+	return types.NewResultOK(nil, "")
+}
+
+func (app *ChainAwareApplication) CheckTx(tx []byte) types.Result {
+	return types.NewResultOK(nil, "")
+}
+
+func (app *ChainAwareApplication) Commit() types.Result {
+	return types.NewResultOK([]byte("nil"), "")
+}
+
+func (app *ChainAwareApplication) Query(query []byte) types.Result {
+	return types.NewResultOK([]byte(Fmt("%d,%d", app.beginCount, app.endCount)), "")
+}
+
+func (app *ChainAwareApplication) BeginBlock(height uint64) {
+	app.beginCount += 1
+	return
+}
+
+func (app *ChainAwareApplication) EndBlock(height uint64) []*types.Validator {
+	app.endCount += 1
+	return nil
+}
+
+func (app *ChainAwareApplication) InitChain(vals []*types.Validator) {
+	return
+}
diff --git a/example/chain_aware/chain_aware_test.go b/example/chain_aware/chain_aware_test.go
new file mode 100644
index 000000000..21eb1b3f9
--- /dev/null
+++ b/example/chain_aware/chain_aware_test.go
@@ -0,0 +1,50 @@
+package main
+
+import (
+	"strconv"
+	"strings"
+	"testing"
+
+	. "github.com/tendermint/go-common"
+	"github.com/tendermint/tmsp/client"
+	"github.com/tendermint/tmsp/server"
+)
+
+func TestChainAware(t *testing.T) {
+
+	app := NewChainAwareApplication()
+
+	// Start the listener
+	_, err := server.NewServer("unix://test.sock", "socket", app)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Connect to the socket
+	client, err := tmspcli.NewSocketClient("unix://test.sock", false)
+	if err != nil {
+		Exit(Fmt("Error starting socket client: %v", err.Error()))
+	}
+	client.Start()
+	defer client.Stop()
+
+	n := uint64(5)
+	for i := uint64(0); i < n; i++ {
+		client.BeginBlockSync(i)
+		client.EndBlockSync(i)
+		client.CommitSync()
+	}
+
+	r := app.Query(nil)
+	spl := strings.Split(string(r.Data), ",")
+	if len(spl) != 2 {
+		t.Fatal("expected %d,%d ; got %s", n, n, string(r.Data))
+	}
+	beginCount, _ := strconv.Atoi(spl[0])
+	endCount, _ := strconv.Atoi(spl[1])
+	if uint64(beginCount) != n {
+		t.Fatalf("expected beginCount of %d, got %d", n, beginCount)
+	} else if uint64(endCount) != n {
+		t.Fatalf("expected endCount of %d, got %d", n, endCount)
+	}
+}
diff --git a/example/dummy/dummy.go b/example/dummy/dummy.go
index ca3a8c03a..21f7c7718 100644
--- a/example/dummy/dummy.go
+++ b/example/dummy/dummy.go
@@ -1,6 +1,7 @@
 package dummy
 
 import (
+	"fmt"
 	"strings"
 
 	. "github.com/tendermint/go-common"
diff --git a/server/socket_server.go b/server/socket_server.go
index 3eec302fb..12d39fefa 100644
--- a/server/socket_server.go
+++ b/server/socket_server.go
@@ -193,6 +193,13 @@ func (s *SocketServer) handleRequest(req *types.Request, responses chan<- *types
 		} else {
 			responses <- types.ToResponseInitChain()
 		}
+	case *types.Request_BeginBlock:
+		if app, ok := s.app.(types.BlockchainAware); ok {
+			app.BeginBlock(r.BeginBlock.Height)
+			responses <- types.ToResponseBeginBlock()
+		} else {
+			responses <- types.ToResponseBeginBlock()
+		}
 	case *types.Request_EndBlock:
 		if app, ok := s.app.(types.BlockchainAware); ok {
 			validators := app.EndBlock(r.EndBlock.Height)