Automatically disconnect agent during integration test cleanup (#228)

internal/integration_tests/agent_tls_termination_test.go

@@ -23,7 +23,6 @@ func TestAgentTLSTerminationIntegration(t *testing.T) {
 
 	// Setup agent
 	agent, ctx := SetupAgent(t)
-	defer func() { _ = agent.Disconnect() }()
 
 	// Setup synchronization primitives
 	handlerReady := testutil.NewSyncPoint()

internal/integration_tests/early_response_test.go

@@ -19,7 +19,6 @@ func TestEarlyResponseLargeUpload(t *testing.T) {
 	t.Parallel()
 
 	agent, ctx := SetupAgent(t)
-	defer func() { _ = agent.Disconnect() }()
 
 	const maxBodySize = 1024
 

internal/integration_tests/endpoint_closing_test.go

@@ -16,7 +16,6 @@ func TestEndpointClosingIntegration(t *testing.T) {
 
 	// Setup agent
 	agent, ctx := SetupAgent(t)
-	defer func() { _ = agent.Disconnect() }()
 
 	// Create a listener endpoint
 	listener, err := agent.Listen(ctx)

internal/integration_tests/forward_test.go

@@ -19,7 +19,6 @@ func TestForward(t *testing.T) {
 	t.Parallel()
 	// Setup agent
 	agent, ctx := SetupAgent(t)
-	defer func() { _ = agent.Disconnect() }()
 
 	// Create a channel to signal when the server is ready
 	serverReady := testutil.NewSyncPoint()

internal/integration_tests/http2_test.go

@@ -25,7 +25,6 @@ func TestUpstreamProtocolHTTP2(t *testing.T) {
 
 		// Setup agent for this test
 		agent, ctx := SetupAgent(t)
-		defer func() { _ = agent.Disconnect() }()
 
 		// Set up a test HTTP/2 server
 		srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -92,7 +91,6 @@ func TestUpstreamProtocolHTTP2(t *testing.T) {
 
 		// Setup agent for this test
 		agent, ctx := SetupAgent(t)
-		defer func() { _ = agent.Disconnect() }()
 
 		// Set up a test HTTP/2 server
 		srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

internal/integration_tests/listen_http_test.go

@@ -16,7 +16,6 @@ func TestListenAndHTTPRequest(t *testing.T) {
 	t.Parallel()
 	// Setup agent
 	agent, ctx := SetupAgent(t)
-	defer func() { _ = agent.Disconnect() }()
 
 	// Setup listener
 	listener := SetupListener(t, agent, ctx)

internal/integration_tests/listen_http_url_test.go

@@ -17,7 +17,6 @@ func TestListenWithHTTPURL(t *testing.T) {
 	t.Parallel()
 	// Setup agent
 	agent, ctx := SetupAgent(t)
-	defer func() { _ = agent.Disconnect() }()
 
 	// Setup listener with HTTP URL
 	httpURL := "http://test-http.ngrok.io"

internal/integration_tests/listen_https_test.go

@@ -17,7 +17,6 @@ func TestListenWithHTTPSURL(t *testing.T) {
 	t.Parallel()
 	// Setup agent
 	agent, ctx := SetupAgent(t)
-	defer func() { _ = agent.Disconnect() }()
 
 	// Setup listener with HTTPS URL
 	httpsURL := "https://test-https.ngrok.io"

internal/integration_tests/listen_tcp_test.go

@@ -17,7 +17,6 @@ func TestListenAndTCPConnection(t *testing.T) {
 	t.Parallel()
 	// Setup agent
 	agent, ctx := SetupAgent(t)
-	defer func() { _ = agent.Disconnect() }()
 
 	// Setup TCP listener using the TCP scheme
 	listener := SetupListener(t, agent, ctx, ngrok.WithURL("tcp://"))

internal/integration_tests/proxy_proto_test.go

@@ -311,7 +311,6 @@ func TestProxyProtoIntegration(t *testing.T) {
 
 			// Setup agent
 			agent, ctx := SetupAgent(t)
-			defer func() { _ = agent.Disconnect() }()
 
 			// Create synchronization points
 			handlerReady := testutil.NewSyncPoint()

internal/integration_tests/test_utils.go

@@ -50,6 +50,11 @@ func SetupAgent(t *testing.T) (ngrok.Agent, context.Context) {
 	// Connect the agent
 	err = agent.Connect(ctx)
 	require.NoError(t, err, "Failed to connect agent")
+	t.Cleanup(func() {
+		if err := agent.Disconnect(); err != nil {
+			t.Error("Agent disconnect:", err)
+		}
+	})
 
 	return agent, ctx
 }

internal/integration_tests/upstream_dialer_test.go

@@ -53,7 +53,6 @@ func TestUpstreamDialer(t *testing.T) {
 
 	// Setup agent
 	agent, ctx := SetupAgent(t)
-	defer func() { _ = agent.Disconnect() }()
 
 	// Create a custom dialer that returns an error and has synchronization
 	customDialer := newErroringDialer()

internal/integration_tests/url_pooling_test.go

@@ -24,7 +24,6 @@ func TestListenWithURLAndPooling(t *testing.T) {
 
 	// Setup agent
 	agent, ctx := SetupAgent(t)
-	defer func() { _ = agent.Disconnect() }()
 
 	// Common URL for both endpoints - IMPORTANT: the exact same string must be used for both listeners
 	sharedURL := "https://test-lb.ngrok.io"

internal/integration_tests/websocket_test.go

@@ -19,7 +19,6 @@ func TestWebSocketUpgrade(t *testing.T) {
 	t.Parallel()
 
 	agent, ctx := SetupAgent(t)
-	defer agent.Disconnect() //nolint:errcheck
 
 	// Start an upstream WebSocket server.
 	server := httptest.NewServer(websocket.Handler(func(ws *websocket.Conn) {