From 942985ce512cb7416a0b5c853261039a1cffef71 Mon Sep 17 00:00:00 2001
From: Asim Aslam <asim@chuhnk.me>
Date: Sat, 7 Jan 2017 14:52:23 +0000
Subject: [PATCH] add service test and shutdown based on context

---
 options.go      | 24 ++++++++++++++++++
 service.go      | 27 ++++++++++++++++----
 service_test.go | 65 +++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 111 insertions(+), 5 deletions(-)
 create mode 100644 service_test.go

diff --git a/options.go b/options.go
index 19f8a3b5..053b934c 100644
--- a/options.go
+++ b/options.go
@@ -28,6 +28,8 @@ type Options struct {
 
 	// Before and After funcs
 	BeforeStart []func() error
+	BeforeStop  []func() error
+	AfterStart  []func() error
 	AfterStop   []func() error
 
 	// Other options for implementations of the interface
@@ -43,6 +45,7 @@ func newOptions(opts ...Option) Options {
 		Server:    server.DefaultServer,
 		Registry:  registry.DefaultRegistry,
 		Transport: transport.DefaultTransport,
+		Context:   context.Background(),
 	}
 
 	for _, o := range opts {
@@ -73,6 +76,15 @@ func Client(c client.Client) Option {
 	}
 }
 
+// Context specifies a context for the service
+// Can be used to signal shutdown of the service
+// Can be used for extra option values
+func Context(ctx context.Context) Option {
+	return func(o *Options) {
+		o.Context = ctx
+	}
+}
+
 func Server(s server.Server) Option {
 	return func(o *Options) {
 		o.Server = s
@@ -204,6 +216,18 @@ func BeforeStart(fn func() error) Option {
 	}
 }
 
+func BeforeStop(fn func() error) Option {
+	return func(o *Options) {
+		o.BeforeStop = append(o.BeforeStop, fn)
+	}
+}
+
+func AfterStart(fn func() error) Option {
+	return func(o *Options) {
+		o.AfterStart = append(o.AfterStart, fn)
+	}
+}
+
 func AfterStop(fn func() error) Option {
 	return func(o *Options) {
 		o.AfterStop = append(o.AfterStop, fn)
diff --git a/service.go b/service.go
index 8241dd4d..f0bf76bd 100644
--- a/service.go
+++ b/service.go
@@ -120,10 +120,24 @@ func (s *service) Start() error {
 		return err
 	}
 
+	for _, fn := range s.opts.AfterStart {
+		if err := fn(); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
 func (s *service) Stop() error {
+	var gerr error
+
+	for _, fn := range s.opts.BeforeStop {
+		if err := fn(); err != nil {
+			gerr = err
+		}
+	}
+
 	if err := s.opts.Server.Deregister(); err != nil {
 		return err
 	}
@@ -132,15 +146,12 @@ func (s *service) Stop() error {
 		return err
 	}
 
-	var gerr error
 	for _, fn := range s.opts.AfterStop {
 		if err := fn(); err != nil {
-			// should we bail if it fails?
-			// other funcs will not be executed
-			// seems wrong
 			gerr = err
 		}
 	}
+
 	return gerr
 }
 
@@ -155,7 +166,13 @@ func (s *service) Run() error {
 
 	ch := make(chan os.Signal, 1)
 	signal.Notify(ch, syscall.SIGTERM, syscall.SIGINT, syscall.SIGKILL)
-	<-ch
+
+	select {
+	// wait on kill signal
+	case <-ch:
+	// wait on context cancel
+	case <-s.opts.Context.Done():
+	}
 
 	// exit reg loop
 	close(ex)
diff --git a/service_test.go b/service_test.go
new file mode 100644
index 00000000..d1101107
--- /dev/null
+++ b/service_test.go
@@ -0,0 +1,65 @@
+package micro
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/micro/go-micro/registry/mock"
+	proto "github.com/micro/go-micro/server/debug/proto"
+
+	"golang.org/x/net/context"
+)
+
+func TestService(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	// cancellation context
+	ctx, cancel := context.WithCancel(context.Background())
+
+	// create service
+	service := NewService(
+		Name("test.service"),
+		Context(ctx),
+		Registry(mock.NewRegistry()),
+		AfterStart(func() error {
+			wg.Done()
+			return nil
+		}),
+	)
+
+	// we can't test service.Init as it parses the command line
+	// service.Init()
+
+	// register handler
+	// do that later
+
+	go func() {
+		// wait for start
+		wg.Wait()
+
+		// test call debug
+		req := service.Client().NewRequest(
+			"test.service",
+			"Debug.Health",
+			new(proto.HealthRequest),
+		)
+
+		rsp := new(proto.HealthResponse)
+
+		err := service.Client().Call(context.TODO(), req, rsp)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if rsp.Status != "ok" {
+			t.Fatalf("service response: %s", rsp.Status)
+		}
+
+		// shutdown the service
+		cancel()
+	}()
+
+	// run service
+	service.Run()
+}