From 156b3f1e3b5d2a311154c2214ab9f360eaa598aa Mon Sep 17 00:00:00 2001
From: Vytenis Darulis <vytenis@uber.com>
Date: Tue, 1 Dec 2020 01:33:44 -0500
Subject: [PATCH] add actual ActivePlacement() impl and test

---
 src/aggregator/client/tcp_client.go      |  4 +-
 src/aggregator/client/tcp_client_test.go | 47 ++++++++++++++++++++----
 2 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/src/aggregator/client/tcp_client.go b/src/aggregator/client/tcp_client.go
index bd75142714..4119e3fef4 100644
--- a/src/aggregator/client/tcp_client.go
+++ b/src/aggregator/client/tcp_client.go
@@ -248,8 +248,8 @@ func (c *TCPClient) ActivePlacement() (placement.Placement, int, error) {
 		return nil, 0, err
 	}
 	defer onPlacementDoneFn()
-	// TODO: expose version once related PR lands
-	return placement.Clone(), 0, nil
+
+	return placement.Clone(), stagedPlacement.Version(), nil
 }
 
 //nolint:gocritic
diff --git a/src/aggregator/client/tcp_client_test.go b/src/aggregator/client/tcp_client_test.go
index 95c0bea3dd..c379aa64fd 100644
--- a/src/aggregator/client/tcp_client_test.go
+++ b/src/aggregator/client/tcp_client_test.go
@@ -29,6 +29,7 @@ import (
 	"time"
 
 	"github.com/golang/mock/gomock"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/m3db/m3/src/cluster/kv/mem"
@@ -197,14 +198,6 @@ var (
 			SetInstances(testPlacementInstances)
 )
 
-func mustNewTestTCPClient(t *testing.T, opts Options) *TCPClient {
-	c, err := NewClient(opts)
-	require.NoError(t, err)
-	value, ok := c.(*TCPClient)
-	require.True(t, ok)
-	return value
-}
-
 func TestTCPClientWriteUntimedMetricClosed(t *testing.T) {
 	c := mustNewTestTCPClient(t, testOptions())
 	require.NoError(t, c.Close())
@@ -794,6 +787,44 @@ func TestTCPClientWriteTimeRangeFor(t *testing.T) {
 	}
 }
 
+func TestTCPClientInitAndClose(t *testing.T) {
+	c := mustNewTestTCPClient(t, testOptions())
+	require.NoError(t, c.Init())
+	require.NoError(t, c.Close())
+}
+
+func TestTCPClientActivePlacement(t *testing.T) {
+	var (
+		c               = mustNewTestTCPClient(t, testOptions())
+		emptyPl         = placement.NewPlacement()
+		ctrl            = gomock.NewController(t)
+		mockPl          = placement.NewMockPlacement(ctrl)
+		stagedPlacement = placement.NewMockActiveStagedPlacement(ctrl)
+		watcher         = placement.NewMockStagedPlacementWatcher(ctrl)
+		doneCalls       int
+	)
+
+	c.placementWatcher = watcher
+	watcher.EXPECT().ActiveStagedPlacement().Return(stagedPlacement, func() { doneCalls++ }, nil)
+	stagedPlacement.EXPECT().Version().Return(42)
+	stagedPlacement.EXPECT().ActivePlacement().Return(mockPl, func() { doneCalls++ }, nil)
+	mockPl.EXPECT().Clone().Return(emptyPl)
+
+	pl, v, err := c.ActivePlacement()
+	assert.NoError(t, err)
+	assert.Equal(t, 42, v)
+	assert.Equal(t, 2, doneCalls)
+	assert.Equal(t, emptyPl, pl)
+}
+
+func mustNewTestTCPClient(t *testing.T, opts Options) *TCPClient {
+	c, err := NewClient(opts)
+	require.NoError(t, err)
+	value, ok := c.(*TCPClient)
+	require.True(t, ok)
+	return value
+}
+
 // TODO: clean this up as it's in use by other test files
 func testOptions() Options {
 	return testTCPClientOptions()