Improve iOS tests

ios/VibeTunnel/Services/BufferWebSocketClient.swift

@@ -49,8 +49,8 @@ class BufferWebSocketClient: NSObject {
     /// Magic byte for binary messages
     private static let bufferMagicByte: UInt8 = 0xBF
 
-    private var webSocketTask: URLSessionWebSocketTask?
-    private let session = URLSession(configuration: .default)
+    private var webSocket: WebSocketProtocol?
+    private let webSocketFactory: WebSocketFactory
     private var subscriptions = [String: (TerminalWebSocketEvent) -> Void]()
     private var reconnectTask: Task<Void, Never>?
     private var reconnectAttempts = 0
@@ -69,6 +69,11 @@ class BufferWebSocketClient: NSObject {
         }
         return serverConfig.baseURL
     }
+    
+    init(webSocketFactory: WebSocketFactory = DefaultWebSocketFactory()) {
+        self.webSocketFactory = webSocketFactory
+        super.init()
+    }
 
     func connect() {
         guard !isConnecting else { return }
@@ -93,39 +98,27 @@ class BufferWebSocketClient: NSObject {
 
         print("[BufferWebSocket] Connecting to \(wsURL)")
 
-        // Cancel existing task if any
-        webSocketTask?.cancel(with: .goingAway, reason: nil)
+        // Disconnect existing WebSocket if any
+        webSocket?.disconnect(with: .goingAway, reason: nil)
 
-        // Create request with authentication
-        var request = URLRequest(url: wsURL)
+        // Create new WebSocket
+        webSocket = webSocketFactory.createWebSocket()
+        webSocket?.delegate = self
 
+        // Build headers
+        var headers: [String: String] = [:]
+        
         // Add authentication header if needed
         if let config = UserDefaults.standard.data(forKey: "savedServerConfig"),
            let serverConfig = try? JSONDecoder().decode(ServerConfig.self, from: config),
            let authHeader = serverConfig.authorizationHeader {
-            request.setValue(authHeader, forHTTPHeaderField: "Authorization")
+            headers["Authorization"] = authHeader
         }
 
-        // Create new WebSocket task
-        webSocketTask = session.webSocketTask(with: request)
-        webSocketTask?.resume()
-
-        // Start receiving messages
-        receiveMessage()
-
-        // Send initial ping to establish connection
+        // Connect
         Task {
             do {
-                try await sendPing()
-                isConnected = true
-                isConnecting = false
-                reconnectAttempts = 0
-                startPingTask()
-
-                // Re-subscribe to all sessions
-                for sessionId in subscriptions.keys {
-                    try await subscribe(to: sessionId)
-                }
+                try await webSocket?.connect(to: wsURL, with: headers)
             } catch {
                 print("[BufferWebSocket] Connection failed: \(error)")
                 connectionError = error
@@ -135,36 +128,13 @@ class BufferWebSocketClient: NSObject {
         }
     }
 
-    private func receiveMessage() {
-        webSocketTask?.receive { [weak self] result in
-            guard let self else { return }
-
-            switch result {
-            case .success(let message):
-                Task { @MainActor in
-                    self.handleMessage(message)
-                    self.receiveMessage() // Continue receiving
-                }
-
-            case .failure(let error):
-                print("[BufferWebSocket] Receive error: \(error)")
-                Task { @MainActor in
-                    self.handleDisconnection()
-                }
-            }
-        }
-    }
-
-    private func handleMessage(_ message: URLSessionWebSocketTask.Message) {
+    private func handleMessage(_ message: WebSocketMessage) {
         switch message {
         case .data(let data):
             handleBinaryMessage(data)
 
         case .string(let text):
             handleTextMessage(text)
-
-        @unknown default:
-            break
         }
     }
 
@@ -632,7 +602,7 @@ class BufferWebSocketClient: NSObject {
     }
 
     private func sendMessage(_ message: [String: Any]) async throws {
-        guard let webSocketTask else {
+        guard let webSocket else {
             throw WebSocketError.connectionFailed
         }
 
@@ -641,11 +611,14 @@ class BufferWebSocketClient: NSObject {
             throw WebSocketError.invalidData
         }
 
-        try await webSocketTask.send(.string(string))
+        try await webSocket.send(.string(string))
     }
 
     private func sendPing() async throws {
-        try await sendMessage(["type": "ping"])
+        guard let webSocket else {
+            throw WebSocketError.connectionFailed
+        }
+        try await webSocket.sendPing()
     }
 
     private func startPingTask() {
@@ -668,7 +641,7 @@ class BufferWebSocketClient: NSObject {
 
     private func handleDisconnection() {
         isConnected = false
-        webSocketTask = nil
+        webSocket = nil
         stopPingTask()
         scheduleReconnect()
     }
@@ -697,8 +670,8 @@ class BufferWebSocketClient: NSObject {
         reconnectTask = nil
         stopPingTask()
 
-        webSocketTask?.cancel(with: .goingAway, reason: nil)
-        webSocketTask = nil
+        webSocket?.disconnect(with: .goingAway, reason: nil)
+        webSocket = nil
 
         subscriptions.removeAll()
         isConnected = false
@@ -706,6 +679,40 @@ class BufferWebSocketClient: NSObject {
 
     deinit {
         // Tasks will be cancelled automatically when the object is deallocated
-        // WebSocket task cleanup happens in disconnect()
+        // WebSocket cleanup happens in disconnect()
+    }
+}
+
+// MARK: - WebSocketDelegate
+
+extension BufferWebSocketClient: WebSocketDelegate {
+    func webSocketDidConnect(_ webSocket: WebSocketProtocol) {
+        print("[BufferWebSocket] Connected")
+        isConnected = true
+        isConnecting = false
+        reconnectAttempts = 0
+        startPingTask()
+        
+        // Re-subscribe to all sessions
+        Task {
+            for sessionId in subscriptions.keys {
+                try? await subscribe(to: sessionId)
+            }
+        }
+    }
+    
+    func webSocket(_ webSocket: WebSocketProtocol, didReceiveMessage message: WebSocketMessage) {
+        handleMessage(message)
+    }
+    
+    func webSocket(_ webSocket: WebSocketProtocol, didFailWithError error: Error) {
+        print("[BufferWebSocket] Error: \(error)")
+        connectionError = error
+        handleDisconnection()
+    }
+    
+    func webSocketDidDisconnect(_ webSocket: WebSocketProtocol, closeCode: URLSessionWebSocketTask.CloseCode, reason: Data?) {
+        print("[BufferWebSocket] Disconnected with code: \(closeCode)")
+        handleDisconnection()
     }
 }

ios/VibeTunnel/Services/WebSocketFactory.swift

@@ -0,0 +1,15 @@
+import Foundation
+
+/// Protocol for creating WebSocket instances
+@MainActor
+protocol WebSocketFactory {
+    func createWebSocket() -> WebSocketProtocol
+}
+
+/// Default factory that creates real WebSocket instances
+@MainActor
+class DefaultWebSocketFactory: WebSocketFactory {
+    func createWebSocket() -> WebSocketProtocol {
+        return URLSessionWebSocket()
+    }
+}

ios/VibeTunnel/Services/WebSocketProtocol.swift

@@ -0,0 +1,152 @@
+import Foundation
+
+/// Protocol for WebSocket operations to enable testing
+@MainActor
+protocol WebSocketProtocol: AnyObject {
+    var delegate: WebSocketDelegate? { get set }
+    
+    func connect(to url: URL, with headers: [String: String]) async throws
+    func send(_ message: WebSocketMessage) async throws
+    func sendPing() async throws
+    func disconnect(with code: URLSessionWebSocketTask.CloseCode, reason: Data?)
+}
+
+/// WebSocket message types
+enum WebSocketMessage {
+    case string(String)
+    case data(Data)
+}
+
+/// Delegate protocol for WebSocket events
+@MainActor
+protocol WebSocketDelegate: AnyObject {
+    func webSocketDidConnect(_ webSocket: WebSocketProtocol)
+    func webSocket(_ webSocket: WebSocketProtocol, didReceiveMessage message: WebSocketMessage)
+    func webSocket(_ webSocket: WebSocketProtocol, didFailWithError error: Error)
+    func webSocketDidDisconnect(_ webSocket: WebSocketProtocol, closeCode: URLSessionWebSocketTask.CloseCode, reason: Data?)
+}
+
+/// Real implementation of WebSocketProtocol using URLSessionWebSocketTask
+@MainActor
+class URLSessionWebSocket: NSObject, WebSocketProtocol {
+    weak var delegate: WebSocketDelegate?
+    private var webSocketTask: URLSessionWebSocketTask?
+    private var session: URLSession!
+    private var isReceiving = false
+    
+    override init() {
+        super.init()
+        self.session = URLSession(configuration: .default, delegate: self, delegateQueue: nil)
+    }
+    
+    func connect(to url: URL, with headers: [String: String]) async throws {
+        var request = URLRequest(url: url)
+        headers.forEach { request.setValue($1, forHTTPHeaderField: $0) }
+        
+        webSocketTask = session.webSocketTask(with: request)
+        webSocketTask?.resume()
+        
+        // Start receiving messages
+        isReceiving = true
+        receiveNextMessage()
+        
+        // Send initial ping to verify connection
+        do {
+            try await sendPing()
+            Task { @MainActor in
+                self.delegate?.webSocketDidConnect(self)
+            }
+        } catch {
+            Task { @MainActor in
+                self.delegate?.webSocket(self, didFailWithError: error)
+            }
+            throw error
+        }
+    }
+    
+    func send(_ message: WebSocketMessage) async throws {
+        guard let task = webSocketTask else {
+            throw WebSocketError.connectionFailed
+        }
+        
+        switch message {
+        case .string(let text):
+            try await task.send(.string(text))
+        case .data(let data):
+            try await task.send(.data(data))
+        }
+    }
+    
+    func sendPing() async throws {
+        guard let task = webSocketTask else {
+            throw WebSocketError.connectionFailed
+        }
+        
+        return try await withCheckedThrowingContinuation { continuation in
+            task.sendPing { error in
+                if let error = error {
+                    continuation.resume(throwing: error)
+                } else {
+                    continuation.resume()
+                }
+            }
+        }
+    }
+    
+    func disconnect(with code: URLSessionWebSocketTask.CloseCode, reason: Data?) {
+        isReceiving = false
+        webSocketTask?.cancel(with: code, reason: reason)
+        Task { @MainActor in
+            self.delegate?.webSocketDidDisconnect(self, closeCode: code, reason: reason)
+        }
+    }
+    
+    private func receiveNextMessage() {
+        guard isReceiving, let task = webSocketTask else { return }
+        
+        task.receive { [weak self] result in
+            guard let self = self else { return }
+            
+            switch result {
+            case .success(let message):
+                let wsMessage: WebSocketMessage
+                switch message {
+                case .string(let text):
+                    wsMessage = .string(text)
+                case .data(let data):
+                    wsMessage = .data(data)
+                @unknown default:
+                    return
+                }
+                
+                Task { @MainActor in
+                    self.delegate?.webSocket(self, didReceiveMessage: wsMessage)
+                }
+                
+                // Continue receiving
+                Task { @MainActor in
+                    self.receiveNextMessage()
+                }
+                
+            case .failure(let error):
+                Task { @MainActor in
+                    self.isReceiving = false
+                    self.delegate?.webSocket(self, didFailWithError: error)
+                }
+            }
+        }
+    }
+}
+
+extension URLSessionWebSocket: URLSessionWebSocketDelegate {
+    nonisolated func urlSession(_ session: URLSession, webSocketTask: URLSessionWebSocketTask, didOpenWithProtocol protocol: String?) {
+        // Connection opened - already handled in connect()
+    }
+    
+    nonisolated func urlSession(_ session: URLSession, webSocketTask: URLSessionWebSocketTask, didCloseWith closeCode: URLSessionWebSocketTask.CloseCode, reason: Data?) {
+        Task { @MainActor in
+            self.isReceiving = false
+            self.delegate?.webSocketDidDisconnect(self, closeCode: closeCode, reason: reason)
+        }
+    }
+}

ios/VibeTunnelTests/Mocks/MockWebSocket.swift

@@ -0,0 +1,196 @@
+import Foundation
+@testable import VibeTunnel
+
+/// Mock WebSocket implementation for testing
+@MainActor
+class MockWebSocket: WebSocketProtocol {
+    weak var delegate: WebSocketDelegate?
+    
+    // State tracking
+    private(set) var isConnected = false
+    private(set) var lastConnectURL: URL?
+    private(set) var lastConnectHeaders: [String: String]?
+    private(set) var sentMessages: [WebSocketMessage] = []
+    private(set) var pingCount = 0
+    private(set) var disconnectCalled = false
+    private(set) var lastDisconnectCode: URLSessionWebSocketTask.CloseCode?
+    private(set) var lastDisconnectReason: Data?
+    
+    // Control test behavior
+    var shouldFailConnection = false
+    var connectionError: Error?
+    var shouldFailSend = false
+    var sendError: Error?
+    var shouldFailPing = false
+    var pingError: Error?
+    
+    // Message simulation
+    private var messageQueue: [WebSocketMessage] = []
+    private var messageDeliveryTask: Task<Void, Never>?
+    
+    func connect(to url: URL, with headers: [String: String]) async throws {
+        lastConnectURL = url
+        lastConnectHeaders = headers
+        
+        if shouldFailConnection {
+            let error = connectionError ?? WebSocketError.connectionFailed
+            throw error
+        }
+        
+        isConnected = true
+        delegate?.webSocketDidConnect(self)
+        
+        // Start delivering queued messages
+        startMessageDelivery()
+    }
+    
+    func send(_ message: WebSocketMessage) async throws {
+        guard isConnected else {
+            throw WebSocketError.connectionFailed
+        }
+        
+        if shouldFailSend {
+            throw sendError ?? WebSocketError.connectionFailed
+        }
+        
+        sentMessages.append(message)
+    }
+    
+    func sendPing() async throws {
+        guard isConnected else {
+            throw WebSocketError.connectionFailed
+        }
+        
+        if shouldFailPing {
+            throw pingError ?? WebSocketError.connectionFailed
+        }
+        
+        pingCount += 1
+    }
+    
+    func disconnect(with code: URLSessionWebSocketTask.CloseCode, reason: Data?) {
+        disconnectCalled = true
+        lastDisconnectCode = code
+        lastDisconnectReason = reason
+        
+        if isConnected {
+            isConnected = false
+            messageDeliveryTask?.cancel()
+            messageDeliveryTask = nil
+            delegate?.webSocketDidDisconnect(self, closeCode: code, reason: reason)
+        }
+    }
+    
+    // MARK: - Test Helpers
+    
+    /// Simulate receiving a message from the server
+    func simulateMessage(_ message: WebSocketMessage) {
+        guard isConnected else { return }
+        messageQueue.append(message)
+    }
+    
+    /// Simulate multiple messages
+    func simulateMessages(_ messages: [WebSocketMessage]) {
+        guard isConnected else { return }
+        messageQueue.append(contentsOf: messages)
+    }
+    
+    /// Simulate a connection error
+    func simulateError(_ error: Error) {
+        guard isConnected else { return }
+        delegate?.webSocket(self, didFailWithError: error)
+    }
+    
+    /// Simulate server disconnection
+    func simulateDisconnection(closeCode: URLSessionWebSocketTask.CloseCode = .abnormalClosure, reason: Data? = nil) {
+        guard isConnected else { return }
+        isConnected = false
+        messageDeliveryTask?.cancel()
+        messageDeliveryTask = nil
+        delegate?.webSocketDidDisconnect(self, closeCode: closeCode, reason: reason)
+    }
+    
+    /// Clear all tracked state
+    func reset() {
+        isConnected = false
+        lastConnectURL = nil
+        lastConnectHeaders = nil
+        sentMessages.removeAll()
+        pingCount = 0
+        disconnectCalled = false
+        lastDisconnectCode = nil
+        lastDisconnectReason = nil
+        messageQueue.removeAll()
+        messageDeliveryTask?.cancel()
+        messageDeliveryTask = nil
+    }
+    
+    /// Find sent messages by type
+    func sentStringMessages() -> [String] {
+        sentMessages.compactMap { message in
+            if case .string(let text) = message {
+                return text
+            }
+            return nil
+        }
+    }
+    
+    func sentDataMessages() -> [Data] {
+        sentMessages.compactMap { message in
+            if case .data(let data) = message {
+                return data
+            }
+            return nil
+        }
+    }
+    
+    /// Find sent JSON messages
+    func sentJSONMessages() -> [[String: Any]] {
+        sentStringMessages().compactMap { string in
+            guard let data = string.data(using: .utf8),
+                  let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
+                return nil
+            }
+            return json
+        }
+    }
+    
+    private func startMessageDelivery() {
+        messageDeliveryTask = Task { [weak self] in
+            while !Task.isCancelled {
+                guard let self = self else { break }
+                
+                if !messageQueue.isEmpty {
+                    let message = messageQueue.removeFirst()
+                    await MainActor.run {
+                        self.delegate?.webSocket(self, didReceiveMessage: message)
+                    }
+                }
+                
+                // Small delay to simulate network latency
+                try? await Task.sleep(nanoseconds: 10_000_000) // 10ms
+            }
+        }
+    }
+}
+
+/// Mock WebSocket factory for testing
+@MainActor
+class MockWebSocketFactory: WebSocketFactory {
+    private(set) var createdWebSockets: [MockWebSocket] = []
+    
+    func createWebSocket() -> WebSocketProtocol {
+        let webSocket = MockWebSocket()
+        createdWebSockets.append(webSocket)
+        return webSocket
+    }
+    
+    var lastCreatedWebSocket: MockWebSocket? {
+        createdWebSockets.last
+    }
+    
+    func reset() {
+        createdWebSockets.forEach { $0.reset() }
+        createdWebSockets.removeAll()
+    }
+}

ios/VibeTunnelTests/Mocks/MockWebSocketTask.swift

@@ -1,45 +1,3 @@
-import Foundation
-
-/// Simple mock WebSocket session for testing
-/// Note: This is a placeholder implementation since we can't easily mock URLSessionWebSocketTask
-/// Real tests should use dependency injection or network stubbing libraries
-class MockWebSocketSession {
-    var mockTask: Any?
-    var lastURL: URL?
-    var lastRequest: URLRequest?
-    
-    func webSocketTask(with url: URL) -> Any {
-        lastURL = url
-        return NSObject() // Return a dummy object
-    }
-    
-    func webSocketTask(with request: URLRequest) -> Any {
-        lastRequest = request
-        return NSObject() // Return a dummy object
-    }
-}
-
-/// Placeholder for future WebSocket testing implementation
-/// Currently, WebSocket tests are limited to conceptual testing
-/// due to URLSessionWebSocketTask not being easily mockable
-struct WebSocketTestHelper {
-    static func createMockBinaryMessage(cols: Int32, rows: Int32) -> Data {
-        var data = Data()
-        
-        // Magic byte
-        data.append(0xBF)
-        
-        // Header (5 Int32 values in little endian)
-        let viewportY: Int32 = 0
-        let cursorX: Int32 = 0
-        let cursorY: Int32 = 0
-        
-        data.append(contentsOf: withUnsafeBytes(of: cols.littleEndian) { Array($0) })
-        data.append(contentsOf: withUnsafeBytes(of: rows.littleEndian) { Array($0) })
-        data.append(contentsOf: withUnsafeBytes(of: viewportY.littleEndian) { Array($0) })
-        data.append(contentsOf: withUnsafeBytes(of: cursorX.littleEndian) { Array($0) })
-        data.append(contentsOf: withUnsafeBytes(of: cursorY.littleEndian) { Array($0) })
-        
-        return data
-    }
-}
+// This file is kept for Xcode project compatibility
+// The actual mock implementation is in MockWebSocket.swift
+import Foundation

ios/VibeTunnelTests/Services/BufferWebSocketClientTests.swift

@@ -8,262 +8,274 @@ struct BufferWebSocketClientTests {
     @Test("Connects successfully with valid configuration")
     func successfulConnection() async throws {
         // Arrange
-        let client = BufferWebSocketClient()
+        let mockFactory = MockWebSocketFactory()
+        let client = BufferWebSocketClient(webSocketFactory: mockFactory)
         saveTestServerConfig()
-
-        // Note: This test would require modifying BufferWebSocketClient to accept a custom URLSession
-        // For now, we'll test the connection logic conceptually
-        // let mockSession = MockWebSocketSession()
-        // let mockTask = mockSession.webSocketTask(with: URL(string: "ws://localhost:8888")!)
-
+        
         // Act
         client.connect()
-
+        
+        // Give it a moment to process
+        try await Task.sleep(nanoseconds: 100_000_000) // 100ms
+        
         // Assert
-        // In a real test, we'd verify the WebSocket connection was established
-        // For now, we verify the client doesn't crash and sets appropriate state
+        #expect(mockFactory.createdWebSockets.count == 1)
+        let mockWebSocket = mockFactory.lastCreatedWebSocket
+        #expect(mockWebSocket?.isConnected == true)
+        #expect(mockWebSocket?.lastConnectURL?.absoluteString.contains("/buffers") == true)
+        #expect(client.isConnected == true)
         #expect(client.connectionError == nil)
     }
-
+    
     @Test("Handles connection failure gracefully")
     func connectionFailure() async throws {
         // Arrange
-        let client = BufferWebSocketClient()
-        // Don't save server config to trigger connection failure
-        UserDefaults.standard.removeObject(forKey: "savedServerConfig")
-
-        // Act
+        let mockFactory = MockWebSocketFactory()
+        let client = BufferWebSocketClient(webSocketFactory: mockFactory)
+        saveTestServerConfig()
+        
+        // Configure mock to fail
         client.connect()
-
+        try await Task.sleep(nanoseconds: 50_000_000) // 50ms
+        
+        let mockWebSocket = mockFactory.lastCreatedWebSocket
+        mockWebSocket?.shouldFailConnection = true
+        mockWebSocket?.connectionError = WebSocketError.connectionFailed
+        
+        // Simulate connection failure
+        mockWebSocket?.simulateError(WebSocketError.connectionFailed)
+        
+        try await Task.sleep(nanoseconds: 50_000_000) // 50ms
+        
         // Assert
-        #expect(client.connectionError != nil)
         #expect(client.isConnected == false)
+        #expect(client.connectionError != nil)
     }
-
+    
     @Test("Parses binary buffer messages correctly")
     func binaryMessageParsing() async throws {
         // Arrange
-        let client = BufferWebSocketClient()
+        let mockFactory = MockWebSocketFactory()
+        let client = BufferWebSocketClient(webSocketFactory: mockFactory)
+        saveTestServerConfig()
+        
         var receivedEvent: TerminalWebSocketEvent?
-
+        let sessionId = "test-session-123"
+        
         // Subscribe to events
-        client.subscribe(id: "test") { event in
+        client.subscribe(to: sessionId) { event in
             receivedEvent = event
         }
-
-        // Create a mock binary message
-        let bufferData = TestFixtures.bufferSnapshot(cols: 80, rows: 24)
-
-        // Act - Simulate receiving a binary message
-        // This would normally come through the WebSocket
-        // We'd need to expose a method for testing or use dependency injection
-
-        // For demonstration, let's test the parsing logic conceptually
-        #expect(bufferData.first == 0xBF) // Magic byte
-
-        // Verify data structure
-        var offset = 1
-        let cols = bufferData.withUnsafeBytes { bytes in
-            bytes.load(fromByteOffset: offset, as: Int32.self).littleEndian
+        
+        // Connect
+        client.connect()
+        try await Task.sleep(nanoseconds: 100_000_000) // 100ms
+        
+        let mockWebSocket = mockFactory.lastCreatedWebSocket
+        #expect(mockWebSocket?.isConnected == true)
+        
+        // Create a binary message with proper structure
+        var messageData = Data()
+        
+        // Magic byte for buffer message
+        messageData.append(0xBF)
+        
+        // Session ID length (4 bytes, little endian)
+        let sessionIdData = sessionId.data(using: .utf8)!
+        var sessionIdLength = UInt32(sessionIdData.count).littleEndian
+        messageData.append(Data(bytes: &sessionIdLength, count: 4))
+        
+        // Session ID
+        messageData.append(sessionIdData)
+        
+        // Buffer data with header
+        messageData.append(TestFixtures.bufferSnapshot(cols: 80, rows: 24))
+        
+        // Act - Simulate receiving the message
+        mockWebSocket?.simulateMessage(.data(messageData))
+        
+        // Wait for processing
+        try await Task.sleep(nanoseconds: 100_000_000) // 100ms
+        
+        // Assert
+        #expect(receivedEvent != nil)
+        if case .bufferUpdate(let snapshot) = receivedEvent {
+            #expect(snapshot.cols == 80)
+            #expect(snapshot.rows == 24)
+        } else {
+            Issue.record("Expected buffer update event")
         }
-        offset += 4
-
-        let rows = bufferData.withUnsafeBytes { bytes in
-            bytes.load(fromByteOffset: offset, as: Int32.self).littleEndian
-        }
-
-        #expect(cols == 80)
-        #expect(rows == 24)
     }
-
+    
     @Test("Handles text messages for events")
     func textMessageHandling() async throws {
         // Arrange
-        let client = BufferWebSocketClient()
-        var receivedEvents: [TerminalWebSocketEvent] = []
-
-        client.subscribe(id: "test") { event in
-            receivedEvents.append(event)
-        }
-
-        // Test various text message formats
-        let messages = [
-            """
-            {"type":"exit","code":0}
-            """,
-            """
-            {"type":"bell"}
-            """,
-            """
-            {"type":"alert","title":"Warning","message":"Session timeout"}
-            """
-        ]
-
-        // Act & Assert
-        // In a real implementation, we'd send these through the WebSocket
-        // and verify the correct events are generated
-
-        // Verify JSON structure
-        for message in messages {
-            let data = message.data(using: .utf8)!
-            let json = try JSONSerialization.jsonObject(with: data) as? [String: Any]
-            #expect(json != nil)
-            #expect(json?["type"] as? String != nil)
-        }
-    }
-
-    @Test("Manages subscriptions correctly")
-    func subscriptionManagement() async throws {
-        // Arrange
-        let client = BufferWebSocketClient()
-        var subscriber1Count = 0
-        var subscriber2Count = 0
-
-        // Act
-        client.subscribe(id: "sub1") { _ in
-            subscriber1Count += 1
-        }
-
-        client.subscribe(id: "sub2") { _ in
-            subscriber2Count += 1
-        }
-
-        // Simulate an event (would normally come through WebSocket)
-        // For testing purposes, we'd need to expose internal methods
-
-        // Remove one subscription
-        client.unsubscribe(id: "sub1")
-
-        // Assert
-        // After unsubscribing sub1, only sub2 should receive events
-        // This would be verified in a full integration test
-    }
-
-    @Test("Handles reconnection with exponential backoff")
-    func reconnectionLogic() async throws {
-        // Arrange
-        let client = BufferWebSocketClient()
+        let mockFactory = MockWebSocketFactory()
+        let client = BufferWebSocketClient(webSocketFactory: mockFactory)
         saveTestServerConfig()
-
-        // Act
+        
+        // Connect
         client.connect()
-
-        // Simulate disconnection
-        // In a real test, we'd trigger this through the WebSocket mock
-
-        // Assert
-        // Verify reconnection attempts happen with increasing delays
-        // This would require exposing reconnection state or using time-based testing
+        try await Task.sleep(nanoseconds: 100_000_000) // 100ms
+        
+        let mockWebSocket = mockFactory.lastCreatedWebSocket
+        #expect(mockWebSocket?.isConnected == true)
+        
+        // Act - Simulate ping message
+        mockWebSocket?.simulateMessage(.string("{\"type\":\"ping\"}"))
+        
+        // Wait for processing
+        try await Task.sleep(nanoseconds: 50_000_000) // 50ms
+        
+        // Assert - Check if pong was sent
+        let sentMessages = mockWebSocket?.sentJSONMessages() ?? []
+        #expect(sentMessages.contains { $0["type"] as? String == "pong" })
     }
-
-    @Test("Cleans up resources on disconnect")
-    func disconnectCleanup() async throws {
+    
+    @Test("Subscribes to sessions correctly")
+    func sessionSubscription() async throws {
         // Arrange
-        let client = BufferWebSocketClient()
+        let mockFactory = MockWebSocketFactory()
+        let client = BufferWebSocketClient(webSocketFactory: mockFactory)
         saveTestServerConfig()
-
+        
+        let sessionId = "test-session-456"
         var eventReceived = false
-        client.subscribe(id: "test") { _ in
+        
+        // Act
+        client.subscribe(to: sessionId) { _ in
             eventReceived = true
         }
-
+        
+        client.connect()
+        try await Task.sleep(nanoseconds: 100_000_000) // 100ms
+        
+        let mockWebSocket = mockFactory.lastCreatedWebSocket
+        
+        // Assert - Check if subscribe message was sent
+        let sentMessages = mockWebSocket?.sentJSONMessages() ?? []
+        #expect(sentMessages.contains { msg in
+            msg["type"] as? String == "subscribe" &&
+            msg["sessionId"] as? String == sessionId
+        })
+    }
+    
+    @Test("Handles reconnection after disconnection")
+    func reconnection() async throws {
+        // Arrange
+        let mockFactory = MockWebSocketFactory()
+        let client = BufferWebSocketClient(webSocketFactory: mockFactory)
+        saveTestServerConfig()
+        
+        // Connect
+        client.connect()
+        try await Task.sleep(nanoseconds: 100_000_000) // 100ms
+        
+        let firstWebSocket = mockFactory.lastCreatedWebSocket
+        #expect(client.isConnected == true)
+        
+        // Act - Simulate disconnection
+        firstWebSocket?.simulateDisconnection()
+        
+        // Wait for reconnection attempt
+        try await Task.sleep(nanoseconds: 2_000_000_000) // 2s
+        
+        // Assert
+        #expect(mockFactory.createdWebSockets.count > 1)
+        let secondWebSocket = mockFactory.lastCreatedWebSocket
+        #expect(secondWebSocket !== firstWebSocket)
+    }
+    
+    @Test("Sends ping messages periodically")
+    func pingMessages() async throws {
+        // Arrange
+        let mockFactory = MockWebSocketFactory()
+        let client = BufferWebSocketClient(webSocketFactory: mockFactory)
+        saveTestServerConfig()
+        
         // Act
         client.connect()
+        try await Task.sleep(nanoseconds: 100_000_000) // 100ms
+        
+        let mockWebSocket = mockFactory.lastCreatedWebSocket
+        let initialPingCount = mockWebSocket?.pingCount ?? 0
+        
+        // Assert - Initial ping during connection
+        #expect(initialPingCount > 0)
+    }
+    
+    @Test("Unsubscribes from sessions correctly")
+    func sessionUnsubscription() async throws {
+        // Arrange
+        let mockFactory = MockWebSocketFactory()
+        let client = BufferWebSocketClient(webSocketFactory: mockFactory)
+        saveTestServerConfig()
+        
+        let sessionId = "test-session-789"
+        
+        // Subscribe first
+        client.subscribe(to: sessionId) { _ in }
+        
+        // Connect
+        client.connect()
+        try await Task.sleep(nanoseconds: 100_000_000) // 100ms
+        
+        let mockWebSocket = mockFactory.lastCreatedWebSocket
+        
+        // Clear sent messages
+        mockWebSocket?.sentMessages.removeAll()
+        
+        // Act - Unsubscribe
+        client.unsubscribe(from: sessionId)
+        try await Task.sleep(nanoseconds: 50_000_000) // 50ms
+        
+        // Assert
+        let sentMessages = mockWebSocket?.sentJSONMessages() ?? []
+        #expect(sentMessages.contains { msg in
+            msg["type"] as? String == "unsubscribe" &&
+            msg["sessionId"] as? String == sessionId
+        })
+    }
+    
+    @Test("Cleans up on disconnect")
+    func cleanup() async throws {
+        // Arrange
+        let mockFactory = MockWebSocketFactory()
+        let client = BufferWebSocketClient(webSocketFactory: mockFactory)
+        saveTestServerConfig()
+        
+        // Subscribe to a session
+        client.subscribe(to: "test-session") { _ in }
+        
+        // Connect
+        client.connect()
+        try await Task.sleep(nanoseconds: 100_000_000) // 100ms
+        
+        let mockWebSocket = mockFactory.lastCreatedWebSocket
+        #expect(client.isConnected == true)
+        
+        // Act
         client.disconnect()
-
+        
         // Assert
         #expect(client.isConnected == false)
-        #expect(client.connectionError == nil)
-
-        // Verify subscriptions are maintained but not receiving events
-        // In a real test, we'd verify no events are delivered after disconnect
-    }
-
-    @Test("Validates magic byte in binary messages")
-    func magicByteValidation() async throws {
-        // Arrange
-        var invalidData = Data([0xAB]) // Wrong magic byte
-        invalidData.append(contentsOf: [0, 0, 0, 0]) // Some dummy data
-
-        // Act & Assert
-        // In the real implementation, this should be rejected
-        #expect(invalidData.first != 0xBF)
-    }
-
-    @Test("Handles malformed JSON gracefully")
-    func malformedJSONHandling() async throws {
-        // Arrange
-        let malformedMessages = [
-            "not json",
-            "{invalid json}",
-            """
-            {"type": }
-            """,
-            ""
-        ]
-
-        // Act & Assert
-        for message in malformedMessages {
-            let data = message.data(using: .utf8) ?? Data()
-            let json = try? JSONSerialization.jsonObject(with: data)
-            #expect(json == nil)
-        }
-    }
-
-    @Test("Maintains connection with periodic pings")
-    func pingMechanism() async throws {
-        // Arrange
-        let client = BufferWebSocketClient()
-        saveTestServerConfig()
-
-        // Act
-        client.connect()
-
-        // Assert
-        // In a real test with mock WebSocket, we'd verify:
-        // 1. Ping messages are sent periodically
-        // 2. Connection stays alive with pings
-        // 3. Connection closes if pings fail
-    }
-
-    // MARK: - Helper Methods
-
-    private func saveTestServerConfig() {
-        let config = TestFixtures.validServerConfig
-        if let data = try? JSONEncoder().encode(config) {
-            UserDefaults.standard.set(data, forKey: "savedServerConfig")
-        }
+        #expect(mockWebSocket?.disconnectCalled == true)
+        #expect(mockWebSocket?.lastDisconnectCode == .goingAway)
     }
 }
 
-// MARK: - Integration Tests
+// MARK: - Test Helpers
 
-@Suite("BufferWebSocketClient Integration Tests", .tags(.integration, .websocket))
-@MainActor
-struct BufferWebSocketClientIntegrationTests {
-    @Test("Full connection and message flow", .timeLimit(.seconds(5)))
-    func fullConnectionFlow() async throws {
-        // This test would require a mock WebSocket server
-        // or modifications to BufferWebSocketClient to accept mock dependencies
-
-        // Arrange
-        let client = BufferWebSocketClient()
-        let expectation = confirmation("Received buffer update")
-
-        client.subscribe(id: "integration-test") { event in
-            if case .bufferUpdate = event {
-                Task { await expectation.fulfill() }
-            }
-        }
-
-        // Act
-        // In a real integration test:
-        // 1. Start mock WebSocket server
-        // 2. Connect client
-        // 3. Send buffer update from server
-        // 4. Verify client receives and parses it correctly
-
-        // Assert
-        // await fulfillment(of: [expectation], timeout: .seconds(2))
+private func saveTestServerConfig() {
+    let config = ServerConfig(
+        host: "localhost",
+        port: 8888,
+        useSSL: false,
+        username: nil,
+        password: nil
+    )
+    
+    if let data = try? JSONEncoder().encode(config) {
+        UserDefaults.standard.set(data, forKey: "savedServerConfig")
     }
-}
+}