implement unit test to verify jetbrains functionality

Emyrk · Emyrk · commit 76d3a24094f8 · 2023-11-30T16:39:42.000-06:00
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -206,12 +206,41 @@ func TestAgent_Stats_Magic(t *testing.T) {
 		remotePort := tcpAddr.Port
 		go echoOnce(t, rl)
 
-		sshClient := setupAgentSSHClient(ctx, t)
+		//nolint:dogsled
+		conn, _, stats, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
+		sshClient, err := conn.SSHClient(ctx)
+		require.NoError(t, err)
 
-		conn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
+		tunneledConn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
 		require.NoError(t, err)
-		defer conn.Close()
-		requireEcho(t, conn)
+		t.Cleanup(func() {
+			// always close on failure of test
+			_ = conn.Close()
+			_ = tunneledConn.Close()
+		})
+
+		var s *agentsdk.Stats
+		require.Eventuallyf(t, func() bool {
+			var ok bool
+			s, ok = <-stats
+			return ok && s.ConnectionCount > 0 &&
+				s.SessionCountJetBrains == 1
+		}, testutil.WaitLong, testutil.IntervalFast,
+			"never saw stats with conn open: %+v", s,
+		)
+
+		// Manually closing the connection
+		requireEcho(t, tunneledConn)
+		_ = rl.Close()
+
+		require.Eventuallyf(t, func() bool {
+			var ok bool
+			s, ok = <-stats
+			return ok && s.ConnectionCount == 0 &&
+				s.SessionCountJetBrains == 0
+		}, testutil.WaitLong, testutil.IntervalFast,
+			"never saw stats after conn closes: %+v", s,
+		)
 	})
 }
 
diff --git a/agent/agentssh/jetbrainstrack.go b/agent/agentssh/jetbrainstrack.go
@@ -1,6 +1,8 @@
 package agentssh
 
 import (
+	"sync"
+
 	"cdr.dev/slog"
 	"go.uber.org/atomic"
 	gossh "golang.org/x/crypto/ssh"
@@ -54,10 +56,13 @@ func (w *ChannelAcceptWatcher) Accept() (gossh.Channel, <-chan *gossh.Request, e
 
 type ChannelOnClose struct {
 	gossh.Channel
+	// once ensures close only decrements the counter once.
+	// Because close can be called multiple times.
+	once sync.Once
 	done func()
 }
 
 func (c *ChannelOnClose) Close() error {
-	c.done()
+	c.once.Do(c.done)
 	return c.Channel.Close()
 }
