Encapsulated and mocked the external dependency

Reason: Pave the way for unit tests

go.mod

@@ -5,6 +5,7 @@ go 1.14
 require (
 	github.com/aws/aws-sdk-go v1.29.15
 	github.com/beanstalkd/go-beanstalk v0.0.0-20200526060843-1cc502ecaf3c
+	github.com/golang/mock v1.4.3
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/prometheus/client_golang v0.9.3
 	github.com/spf13/cobra v0.0.5

go.sum

@@ -87,7 +87,10 @@ github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfU
 github.com/golang/groupcache v0.0.0-20160516000752-02826c3e7903/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
 github.com/golang/groupcache v0.0.0-20190129154638-5b532d6fd5ef h1:veQD95Isof8w9/WXiA+pa3tz3fJXkt5B7QaRBrM62gk=
 github.com/golang/groupcache v0.0.0-20190129154638-5b532d6fd5ef/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
+github.com/golang/mock v1.1.1 h1:G5FRp8JnTd7RQH5kemVNlMeyXQAztQ3mOWV95KxsXH8=
 github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A=
+github.com/golang/mock v1.4.3 h1:GV+pQPG/EUUbkh47niozDcADz6go/dUwhVzdUQHIVRw=
+github.com/golang/mock v1.4.3/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt3cw=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.1 h1:YF8+flBXS5eO826T4nzqPrxfhQThhXl0YzfuUPu4SBg=
 github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
@@ -312,6 +315,10 @@ modernc.org/golex v1.0.0/go.mod h1:b/QX9oBD/LhixY6NDh+IdGv17hgB+51fET1i2kPSmvk=
 modernc.org/mathutil v1.0.0/go.mod h1:wU0vUrJsVWBZ4P6e7xtFJEhFSNsfRLJ8H458uRjg03k=
 modernc.org/strutil v1.0.0/go.mod h1:lstksw84oURvj9y3tn8lGvRxyRC1S2+g5uuIzNfIOBs=
 modernc.org/xc v1.0.0/go.mod h1:mRNCo0bvLjGhHO9WsyuKVU4q0ceiDDDoEeWDJHrNx8I=
+rsc.io/quote/v3 v3.1.0 h1:9JKUTTIUgS6kzR9mK1YuGKv6Nl+DijDNIc0ghT58FaY=
+rsc.io/quote/v3 v3.1.0/go.mod h1:yEA65RcK8LyAZtP9Kv3t0HmxON59tX3rD+tICJqUlj0=
+rsc.io/sampler v1.3.0 h1:7uVkIFmeBqHfdjD+gZwtXXI+RODJ2Wc4O7MPEh/QiW4=
+rsc.io/sampler v1.3.0/go.mod h1:T1hPZKmBbMNahiBKFy5HrXp6adAjACjK9JXDnKaTXpA=
 sigs.k8s.io/structured-merge-diff v0.0.0-20190302045857-e85c7b244fd2/go.mod h1:wWxsB5ozmmv/SG7nM11ayaAW51xMvak/t1r0CSlcokI=
 sigs.k8s.io/yaml v1.1.0 h1:4A07+ZFc2wgJwo8YNlQpr1rVlgUDlxXHhPJciaPY5gs=
 sigs.k8s.io/yaml v1.1.0/go.mod h1:UJmg0vDUVViEyp3mgSv9WPwZCDxu4rQW1olrI1uml+o=

pkg/queue/beanstalk.go

@@ -1,6 +1,7 @@
 package queue
 
 import (
+	"errors"
 	"net/url"
 	"path"
 	"strconv"
@@ -14,8 +15,8 @@ import (
 // Beanstalk is used to by the Poller to get the queue
 // information from Beanstalk, it implements the QueuingService interface
 type Beanstalk struct {
-	queues   *Queues
-	connPool *sync.Map
+	queues     *Queues
+	clientPool *sync.Map
 
 	shortPollInterval time.Duration
 	longPollInterval  int64
@@ -27,36 +28,41 @@ func NewBeanstalk(
 	longPollInterval int) (QueuingService, error) {
 
 	return &Beanstalk{
-		queues:   queues,
-		connPool: new(sync.Map),
+		queues:     queues,
+		clientPool: new(sync.Map),
 
 		shortPollInterval: time.Second * time.Duration(shortPollInterval),
 		longPollInterval:  int64(longPollInterval),
 	}, nil
 }
 
-func MustParseInt(s string, base int, bitSize int) int32 {
+func mustParseInt(s string, base int, bitSize int) int32 {
 	i, err := strconv.ParseInt(s, base, bitSize)
 	if err != nil {
 		klog.Fatalf("Error parsing int: %v", err)
 	}
 	return int32(i)
 }
 
-func MustParseUint(s string, base int, bitSize int) uint32 {
+func mustParseUint(s string, base int, bitSize int) uint32 {
 	i, err := strconv.ParseUint(s, base, bitSize)
 	if err != nil {
 		klog.Fatalf("Error parsing int: %v", err)
 	}
 	return uint32(i)
 }
 
-func (b *Beanstalk) getConnection(queueURI string) (*beanstalk.Conn, error) {
-	conn, _ := b.connPool.Load(queueURI)
-	if conn != nil {
-		return conn.(*beanstalk.Conn), nil
-	}
+type BeanstalkClientInterface interface {
+	getStats() (int32, int32, int32, error)
+	longPollReceiveMessage(longPollInterval int64) (int32, int32, error)
+}
+
+type beanstalkClient struct {
+	conn     *beanstalk.Conn
+	queueURI string
+}
 
+func NewBeanstalkClient(queueURI string) (BeanstalkClientInterface, error) {
 	var host, port string
 	parsedURI, err := url.Parse(queueURI)
 	if err != nil {
@@ -69,87 +75,93 @@ func (b *Beanstalk) getConnection(queueURI string) (*beanstalk.Conn, error) {
 		port = "11300"
 	}
 
-	conn, err = beanstalk.Dial("tcp", host+":"+port)
-	if err != nil {
-		return nil, err
-	}
-
-	b.connPool.Store(queueURI, conn)
-
-	return conn.(*beanstalk.Conn), nil
-}
-
-func (b *Beanstalk) getTube(queueURI string) (*beanstalk.Tube, error) {
-	conn, err := b.getConnection(queueURI)
+	conn, err := beanstalk.Dial("tcp", host+":"+port)
 	if err != nil {
-		return nil, err
+		return nil, errors.New("beanstalk dial error: " + err.Error())
 	}
 
-	return &beanstalk.Tube{Conn: conn, Name: path.Base(queueURI)}, nil
+	return &beanstalkClient{conn: conn, queueURI: queueURI}, nil
 }
 
-func (b *Beanstalk) getStats(queueURI string) (int32, int32, int32, error) {
-	tube, err := b.getTube(queueURI)
-	if err != nil {
-		return 0, 0, 0, err
-	}
+func (c *beanstalkClient) getStats() (int32, int32, int32, error) {
+	tube := &beanstalk.Tube{Conn: c.conn, Name: path.Base(c.queueURI)}
 
 	output, err := tube.Stats()
 	if err != nil {
-		return 0, 0, 0, err
+		return 0, 0, 0, errors.New("beanstalk get-stats error: " + err.Error())
 	}
 
-	jobsWaiting := MustParseInt(output["current-jobs-ready"], 10, 32)
-	idleWorkers := MustParseInt(output["current-waiting"], 10, 32)
-	jobsReserved := MustParseInt(output["current-jobs-reserved"], 10, 32)
+	jobsWaiting := mustParseInt(output["current-jobs-ready"], 10, 32)
+	idleWorkers := mustParseInt(output["current-waiting"], 10, 32)
+	jobsReserved := mustParseInt(output["current-jobs-reserved"], 10, 32)
 	return jobsWaiting, idleWorkers, jobsReserved, nil
 }
 
-func (b *Beanstalk) longPollReceiveMessage(queueURI string) (int32, int32, error) {
-	conn, err := b.getConnection(queueURI)
-	if err != nil {
-		return 0, 0, err
-	}
+func (c *beanstalkClient) longPollReceiveMessage(
+	longPollInterval int64) (int32, int32, error) {
 
-	tubeSet := beanstalk.NewTubeSet(conn, path.Base(queueURI))
+	tubeSet := beanstalk.NewTubeSet(c.conn, path.Base(c.queueURI))
 	id, _, err := tubeSet.Reserve(
-		time.Duration(b.longPollInterval) * time.Second,
+		time.Duration(longPollInterval) * time.Second,
 	)
 	e, ok := err.(beanstalk.ConnError)
 	if ok && e.Err == beanstalk.ErrTimeout {
 		return 0, 0, nil
 	}
 	if err != nil {
-		return 0, 0, err
+		return 0, 0, errors.New("beanstalk tube-reserve error: " + err.Error())
 	}
 
-	statsJob, err := conn.StatsJob(id)
+	statsJob, err := c.conn.StatsJob(id)
 	if err != nil {
-		return 0, 0, err
+		return 0, 0, errors.New("beanstalk stats-job error: " + err.Error())
 	}
 
-	conn.Release(id, MustParseUint(statsJob["pri"], 10, 32), 0)
+	c.conn.Release(id, mustParseUint(statsJob["pri"], 10, 32), 0)
 
 	return 1, 0, nil
 }
 
-// TODO: need to get this data from some source
-// like: Prometheus: https://github.com/practo/beanstalkd_exporter
-func (b *Beanstalk) getAverageNumberOfMessagesSent(queueURI string) (float64, error) {
-	return 0.0, nil
+func (b *Beanstalk) getClient(
+	queueURI string) (BeanstalkClientInterface, error) {
+
+	client, _ := b.clientPool.Load(queueURI)
+	if client != nil {
+		return client.(BeanstalkClientInterface), nil
+	}
+
+	client, err := NewBeanstalkClient(queueURI)
+	if err != nil {
+		return nil, err
+	}
+	b.clientPool.Store(queueURI, client)
+
+	return client.(BeanstalkClientInterface), nil
 }
 
 func (b *Beanstalk) getApproxMessages(queueURI string) (int32, error) {
-	jobsWaiting, _, _, err := b.getStats(queueURI)
+	client, err := b.getClient(queueURI)
+	if err != nil {
+		return 0, err
+	}
+
+	jobsWaiting, _, _, err := client.getStats()
 	if err != nil {
 		return jobsWaiting, err
 	}
 
 	return jobsWaiting, nil
 }
 
-func (b *Beanstalk) getApproxMessagesNotVisible(queueURI string) (int32, error) {
-	_, _, jobsReserved, err := b.getStats(queueURI)
+func (b *Beanstalk) getApproxMessagesNotVisible(
+	queueURI string) (int32, error) {
+
+	client, err := b.getClient(queueURI)
+	if err != nil {
+		return 0, err
+	}
+
+	_, _, jobsReserved, err := client.getStats()
 	if err != nil {
 		return jobsReserved, err
 	}
@@ -158,14 +170,41 @@ func (b *Beanstalk) getApproxMessagesNotVisible(queueURI string) (int32, error)
 }
 
 func (b *Beanstalk) getIdleWorkers(queueURI string) (int32, error) {
-	_, idleWorkers, _, err := b.getStats(queueURI)
+	client, err := b.getClient(queueURI)
+	if err != nil {
+		return 0, err
+	}
+
+	_, idleWorkers, _, err := client.getStats()
 	if err != nil {
 		return idleWorkers, err
 	}
 
 	return idleWorkers, nil
 }
 
+// TODO: need to get this data from some source
+// like: Prometheus: https://github.com/practo/beanstalkd_exporter
+func (b *Beanstalk) getAverageNumberOfMessagesSent(
+	queueURI string) (float64, error) {
+
+	return 0.0, nil
+}
+
+func (b *Beanstalk) longPollReceiveMessage(
+	queueURI string) (int32, int32, error) {
+
+	client, err := b.getClient(queueURI)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	messages, idleWorkers, err := client.longPollReceiveMessage(
+		b.longPollInterval)
+
+	return messages, idleWorkers, err
+}
+
 func (b *Beanstalk) Sync(stopCh <-chan struct{}) {
 	for {
 		select {

pkg/queue/beanstalk_test.go

@@ -0,0 +1,21 @@
+package queue
+
+import (
+	_ "github.com/golang/mock/gomock"
+	"testing"
+)
+
+const (
+	queueNameOne = "q1"
+)
+
+func newTestQueues() *Queues {
+	q := NewQueues()
+
+	q.Add(
+		"namespace1",
+		"queuename1",
+		"beanstalk://beanstalkd.namespace1:11300/"+queueNameOne,
+	)
+	return q
+}