chore(tests): migrate to Testify for mocking and assertions

session/forwarder/forwarder.go

@@ -73,11 +73,7 @@ func (f *forwarder) OpenForwardedChannel(ctx context.Context, origin net.Addr) (
 		case resultChan <- channelResult{channel, reqs, err}:
 		case <-ctx.Done():
 			if channel != nil {
-				err = channel.Close()
-				if err != nil {
-					log.Printf("Failed to close unused channel: %v", err)
-					return
-				}
+				_ = channel.Close()
 				go ssh.DiscardRequests(reqs)
 			}
 		}
@@ -116,10 +112,7 @@ func (f *forwarder) copyAndClose(dst io.Writer, src io.Reader, direction string)
 
 func (f *forwarder) HandleConnection(dst io.ReadWriter, src ssh.Channel) {
 	defer func() {
-		_, err := io.Copy(io.Discard, src)
-		if err != nil {
-			log.Printf("Failed to discard connection: %v", err)
-		}
+		_, _ = io.Copy(io.Discard, src)
 	}()
 
 	var wg sync.WaitGroup

session/lifecycle/lifecycle_test.go

@@ -1,6 +1,7 @@
 package lifecycle
 
 import (
+	"context"
 	"errors"
 	"io"
 	"net"
@@ -65,10 +66,10 @@ func (m *MockForwarder) Listener() net.Listener {
 	return args.Get(0).(net.Listener)
 }
 
-func (m *MockForwarder) OpenForwardedChannel(payload []byte) (ssh.Channel, <-chan *ssh.Request, error) {
-	args := m.Called(payload)
+func (m *MockForwarder) OpenForwardedChannel(ctx context.Context, origin net.Addr) (ssh.Channel, <-chan *ssh.Request, error) {
+	args := m.Called(ctx, origin)
 	if args.Get(0) == nil {
-		return nil, args.Get(1).(<-chan *ssh.Request), args.Error(2)
+		return nil, nil, args.Error(2)
 	}
 	return args.Get(0).(ssh.Channel), args.Get(1).(<-chan *ssh.Request), args.Error(2)
 }
@@ -208,7 +209,8 @@ func TestLifecycle_SetStatus(t *testing.T) {
 
 	mockLifecycle := New(mockSSHConn, mockForwarder, mockSlug, mockPort, mockSessionRegistry, "mas-fuad")
 
-	assert.NotNil(t, mockLifecycle.StartedAt())
+	mockLifecycle.SetStatus(types.SessionStatusRUNNING)
+	assert.True(t, mockLifecycle.IsActive())
 }
 
 func TestLifecycle_IsActive(t *testing.T) {