1
1
package agent_test
2
2
3
3
import (
4
+ "bufio"
4
5
"bytes"
5
6
"context"
6
7
"encoding/json"
@@ -152,7 +153,7 @@ func TestAgent_Stats_Magic(t *testing.T) {
152
153
require .NoError (t ,err )
153
154
require .Equal (t ,expected ,strings .TrimSpace (string (output )))
154
155
})
155
- t .Run ("Tracks " ,func (t * testing.T ) {
156
+ t .Run ("TracksVSCode " ,func (t * testing.T ) {
156
157
t .Parallel ()
157
158
if runtime .GOOS == "window" {
158
159
t .Skip ("Sleeping for infinity doesn't work on Windows" )
@@ -192,36 +193,45 @@ func TestAgent_Stats_Magic(t *testing.T) {
192
193
require .NoError (t ,err )
193
194
})
194
195
195
- // This test name must contain the string checked for by the agent, since it
196
- // looks for this string in the process name.
197
- //
198
- // This test sets up a port forward that emulates what Jetbrains IDE's do when
199
- // using gateway. The remote server side of the port forward is spun up using
200
- // the gotest process, which includes the test name.
201
- // So this unit test emulates a PID in the workspace with a similar
202
- // name to the jetbrains IDE. That makes the agent this this SSH port
203
- // forward is a "jetbrains" session.
204
- t .Run ("TracksIdea.vendor.name=JetBrains" ,func (t * testing.T ) {
196
+ t .Run ("TracksJetBrains" ,func (t * testing.T ) {
205
197
t .Parallel ()
206
198
if runtime .GOOS != "linux" {
207
199
t .Skip ("JetBrains tracking is only supported on Linux" )
208
200
}
201
+
209
202
ctx := testutil .Context (t ,testutil .WaitLong )
210
203
211
- rl ,err := net .Listen ("tcp" ,"127.0.0.1:0" )
204
+ // JetBrains tracking works by looking at the process name listening on the
205
+ // forwarded port. If the process's command line includes the magic string
206
+ // we are looking for, then we assume it is a JetBrains editor. So when we
207
+ // connect to the port we must ensure the process includes that magic string
208
+ // to fool the agent into thinking this is JetBrains. To do this we need to
209
+ // spawn an external process (in this case a simple echo server) so we can
210
+ // control the process name. The -D here is just to mimic how Java options
211
+ // are set but is not necessary as the agent looks only for the magic
212
+ // string itself anywhere in the command.
213
+ _ ,b ,_ ,ok := runtime .Caller (0 )
214
+ require .True (t ,ok )
215
+ dir := filepath .Join (filepath .Dir (b ),"../scripts/echoserver/main.go" )
216
+ echoServerCmd := exec .Command ("go" ,"run" ,dir ,
217
+ "-D" ,agentssh .MagicProcessCmdlineJetBrains )
218
+ stdout ,err := echoServerCmd .StdoutPipe ()
212
219
require .NoError (t ,err )
213
- defer rl .Close ()
214
- tcpAddr ,valid := rl .Addr ().(* net.TCPAddr )
215
- require .True (t ,valid )
216
- remotePort := tcpAddr .Port
217
- go echoOnce (t ,rl )
220
+ err = echoServerCmd .Start ()
221
+ require .NoError (t ,err )
222
+ defer echoServerCmd .Process .Kill ()
223
+
224
+ // The echo server prints its port as the first line.
225
+ sc := bufio .NewScanner (stdout )
226
+ sc .Scan ()
227
+ remotePort := sc .Text ()
218
228
219
229
//nolint:dogsled
220
230
conn ,_ ,stats ,_ ,_ := setupAgent (t , agentsdk.Manifest {},0 )
221
231
sshClient ,err := conn .SSHClient (ctx )
222
232
require .NoError (t ,err )
223
233
224
- tunneledConn ,err := sshClient .Dial ("tcp" ,fmt .Sprintf ("127.0.0.1:%d " ,remotePort ))
234
+ tunneledConn ,err := sshClient .Dial ("tcp" ,fmt .Sprintf ("127.0.0.1:%s " ,remotePort ))
225
235
require .NoError (t ,err )
226
236
t .Cleanup (func () {
227
237
// always close on failure of test
@@ -239,9 +249,10 @@ func TestAgent_Stats_Magic(t *testing.T) {
239
249
"never saw stats with conn open: %+v" ,s ,
240
250
)
241
251
242
- //Manually closing the connection
252
+ //Kill theserver and connection after checking for the echo.
243
253
requireEcho (t ,tunneledConn )
244
- _ = rl .Close ()
254
+ _ = echoServerCmd .Process .Kill ()
255
+ _ = tunneledConn .Close ()
245
256
246
257
require .Eventuallyf (t ,func ()bool {
247
258
var ok bool