Add unit test to confirm tracking

Emyrk · Emyrk · commit 244797063de8 · 2023-11-30T16:20:55.000-06:00
diff --git a/agent/agent_test.go b/agent/agent_test.go
@@ -191,6 +191,28 @@ func TestAgent_Stats_Magic(t *testing.T) {
 		err = session.Wait()
 		require.NoError(t, err)
 	})
+
+	// This test name being "Jetbrains" is required to be a certain string.
+	// It must match the regex check in the agent for Jetbrains.
+	t.Run("Jetbrains", func(t *testing.T) {
+		t.Parallel()
+		ctx := testutil.Context(t, testutil.WaitLong)
+
+		rl, err := net.Listen("tcp", "127.0.0.1:0")
+		require.NoError(t, err)
+		defer rl.Close()
+		tcpAddr, valid := rl.Addr().(*net.TCPAddr)
+		require.True(t, valid)
+		remotePort := tcpAddr.Port
+		go echoOnce(t, rl)
+
+		sshClient := setupAgentSSHClient(ctx, t)
+
+		conn, err := sshClient.Dial("tcp", fmt.Sprintf("127.0.0.1:%d", remotePort))
+		require.NoError(t, err)
+		defer conn.Close()
+		requireEcho(t, conn)
+	})
 }
 
 func TestAgent_SessionExec(t *testing.T) {
diff --git a/agent/agentssh/agentssh.go b/agent/agentssh/agentssh.go
@@ -295,8 +295,8 @@ func (s *Server) sessionStart(logger slog.Logger, session ssh.Session, extraEnv
 		s.connCountVSCode.Add(1)
 		defer s.connCountVSCode.Add(-1)
 	case MagicSessionTypeJetBrains:
-		s.connCountJetBrains.Add(1)
-		defer s.connCountJetBrains.Add(-1)
+		// Do nothing here because jetbrains launches hundreds of ssh sessions.
+		// We instead track jetbrains in the single persistent tcp forwarding channel.
 	case "":
 		s.connCountSSHSession.Add(1)
 		defer s.connCountSSHSession.Add(-1)
