From 5c7af5a3babf5b16c6e62952784f518888930d77 Mon Sep 17 00:00:00 2001
From: Joshua Carroll <josh@qntfy.com>
Date: Fri, 1 Feb 2019 00:29:26 +0000
Subject: [PATCH] Add test coverage and handling for Close() with nil Source or
 Sink

---
 frizzle.go      | 15 +++++++++++----
 frizzle_test.go |  8 ++++++++
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/frizzle.go b/frizzle.go
index ca4355f..06f328c 100644
--- a/frizzle.go
+++ b/frizzle.go
@@ -153,6 +153,9 @@ func (f *Friz) Fail(m Msg) error {
 // the configured Source and Sink. Any Msgs which are still unAcked after the
 // timeout has expired are Failed.
 func (f *Friz) FlushAndClose(timeout time.Duration) error {
+	if f.source == nil {
+		panic("FlushAndClose() should not be called without a Source configured; use Close()")
+	}
 	f.source.Stop()
 	tick := time.NewTicker(timeout / 10).C
 	timeoutAlarm := time.After(timeout)
@@ -179,11 +182,15 @@ func (f *Friz) FlushAndClose(timeout time.Duration) error {
 // The Frizzle must not be used afterward.
 func (f *Friz) Close() error {
 	f.log.Debug("Attempting to Close frizzle")
-	if err := f.sink.Close(); err != nil {
-		return err
+	if f.sink != nil {
+		if err := f.sink.Close(); err != nil {
+			return err
+		}
 	}
-	if err := f.source.Close(); err != nil {
-		return err
+	if f.source != nil {
+		if err := f.source.Close(); err != nil {
+			return err
+		}
 	}
 	if f.failSink != nil {
 		if err := f.failSink.Close(); err != nil {
diff --git a/frizzle_test.go b/frizzle_test.go
index 41d55e6..44dcf46 100644
--- a/frizzle_test.go
+++ b/frizzle_test.go
@@ -51,6 +51,10 @@ func TestNilSource(t *testing.T) {
 
 	err := f.Send(testMsg, testSinkName)
 	assert.Nil(t, err)
+
+	mSink.On("Close").Return(nil)
+	err = f.Close()
+	assert.Nil(t, err)
 	mSink.AssertExpectations(t)
 }
 
@@ -80,6 +84,10 @@ func TestNilSink(t *testing.T) {
 
 	receivedMsg := <-f.Receive()
 	assert.Equal(t, testMsgData, receivedMsg.Data())
+
+	mSource.On("Close").Return(nil)
+	err := f.Close()
+	assert.Nil(t, err)
 	mSource.AssertExpectations(t)
 }