fix WebSocket close logic & tests

ronnyf · Jan 1, 2024 · 45550c1 · 45550c1
1 parent 3c0e5ca
commit 45550c1
Show file tree

Hide file tree

Showing 15 changed files with 42 additions and 36 deletions.
diff --git a/Sources/LiveKit/Core/Engine+SignalClientDelegate.swift b/Sources/LiveKit/Core/Engine+SignalClientDelegate.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Core/Engine+TransportDelegate.swift b/Sources/LiveKit/Core/Engine+TransportDelegate.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Core/Engine.swift b/Sources/LiveKit/Core/Engine.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Core/Room+SignalClientDelegate.swift b/Sources/LiveKit/Core/Room+SignalClientDelegate.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Core/Room.swift b/Sources/LiveKit/Core/Room.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Core/SignalClient.swift b/Sources/LiveKit/Core/SignalClient.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Core/Transport.swift b/Sources/LiveKit/Core/Transport.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Extensions/CustomStringConvertible.swift b/Sources/LiveKit/Extensions/CustomStringConvertible.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Protocols/SignalClientDelegate.swift b/Sources/LiveKit/Protocols/SignalClientDelegate.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Support/AsyncCompleter.swift b/Sources/LiveKit/Support/AsyncCompleter.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Support/AsyncQueueActor.swift b/Sources/LiveKit/Support/AsyncQueueActor.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

diff --git a/Sources/LiveKit/Support/WebSocket.swift b/Sources/LiveKit/Support/WebSocket.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -68,10 +68,10 @@ class WebSocket: NSObject, Loggable, AsyncSequence, URLSessionWebSocketDelegate
     }
 
     func close() {
-        task.cancel(with: .goingAway, reason: nil)
+        task.cancel(with: .normalClosure, reason: nil)
         connectContinuation?.resume(throwing: LiveKitError(.cancelled))
         connectContinuation = nil
-        streamContinuation?.finish()
+        streamContinuation?.finish(throwing: LiveKitError(.cancelled))
         streamContinuation = nil
     }
 
@@ -83,7 +83,7 @@ class WebSocket: NSObject, Loggable, AsyncSequence, URLSessionWebSocketDelegate
 
     private func waitForNextValue() {
         guard task.closeCode == .invalid else {
-            streamContinuation?.finish()
+            streamContinuation?.finish(throwing: LiveKitError(.invalidState))
             streamContinuation = nil
             return
         }
@@ -98,7 +98,7 @@ class WebSocket: NSObject, Loggable, AsyncSequence, URLSessionWebSocketDelegate
                 continuation.yield(message)
                 self?.waitForNextValue()
             } catch {
-                continuation.finish(throwing: error)
+                continuation.finish(throwing: LiveKitError.from(error: error))
                 self?.streamContinuation = nil
             }
         })
@@ -119,10 +119,18 @@ class WebSocket: NSObject, Loggable, AsyncSequence, URLSessionWebSocketDelegate
     }
 
     func urlSession(_: URLSession, task _: URLSessionTask, didCompleteWithError error: Error?) {
-        log("didCompleteWithError: \(String(describing: error))", .error)
-        connectContinuation?.resume(throwing: LiveKitError.from(error: error) ?? LiveKitError(.unknown))
+        log("didCompleteWithError: \(String(describing: error))", error != nil ? .error : .debug)
+
+        if let error {
+            let lkError = LiveKitError.from(error: error) ?? LiveKitError(.unknown)
+            connectContinuation?.resume(throwing: lkError)
+            streamContinuation?.finish(throwing: lkError)
+        } else {
+            connectContinuation?.resume()
+            streamContinuation?.finish()
+        }
+
         connectContinuation = nil
-        streamContinuation?.finish()
         streamContinuation = nil
     }
 }
diff --git a/Tests/LiveKitTests/AsyncRetryTests.swift b/Tests/LiveKitTests/AsyncRetryTests.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -25,7 +25,7 @@ class AsyncRetryTests: XCTestCase {
     func testRetry1() async throws {
         let test = Task.retrying(maxRetryCount: 1) { totalAttempts, currentAttempt in
             print("[TEST] Retrying with remaining attemps: \(currentAttempt)/\(totalAttempts)...")
-            throw EngineError.state(message: "Test error")
+            throw LiveKitError(.invalidState, message: "Test error")
             return "Complete"
         }
 

diff --git a/Tests/LiveKitTests/CompleterTests.swift b/Tests/LiveKitTests/CompleterTests.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,13 +26,13 @@ class CompleterTests: XCTestCase {
         let completer = AsyncCompleter<Void>(label: "Test01", timeOut: .seconds(1))
         do {
             try await completer.wait()
-        } catch AsyncCompleterError.timedOut {
+        } catch let error as LiveKitError where error.type == .timedOut {
             print("Timed out 1")
         }
         // Re-use
         do {
             try await completer.wait()
-        } catch AsyncCompleterError.timedOut {
+        } catch let error as LiveKitError where error.type == .timedOut {
             print("Timed out 2")
         }
     }
@@ -53,14 +53,14 @@ class CompleterTests: XCTestCase {
                     // Cancel after 1 second
                     try await Task.sleep(until: .now + .seconds(1), clock: .continuous)
                     print("Task 2: Cancelling completer...")
-                    completer.cancel()
+                    completer.reset()
                 }
 
                 try await group.waitForAll()
             }
-        } catch let error as AsyncCompleterError where error == .timedOut {
+        } catch let error as LiveKitError where error.type == .timedOut {
             print("Completer timed out")
-        } catch let error as AsyncCompleterError where error == .cancelled {
+        } catch let error as LiveKitError where error.type == .cancelled {
             print("Completer cancelled")
         } catch {
             print("Unknown error: \(error)")

diff --git a/Tests/LiveKitTests/WebSocketTests.swift b/Tests/LiveKitTests/WebSocketTests.swift
@@ -1,5 +1,5 @@
 /*
- * Copyright 2023 LiveKit
+ * Copyright 2024 LiveKit
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -26,22 +26,20 @@ class WebSocketTests: XCTestCase {
         print("Connecting...")
         let socket = try await WebSocket(url: URL(string: "wss://socketsbay.com/wss/v2/1/demo/")!)
 
-        print("Connected. Waiting for messages...")
+        print("Connected, waiting for messages...")
         do {
             for try await message in socket {
                 switch message {
-                case let .string(string):
-                    print(string)
-                case let .data(data):
-                    print(data)
-                @unknown default:
-                    print("unknown message received")
+                case let .string(string): print("Received String: \(string)")
+                case let .data(data): print("Received Data: \(data)")
+                @unknown default: print("Received unknown message")
                 }
             }
         } catch {
             print("Error: \(error)")
+            throw error
         }
 
-        print("Completed.")
+        print("Completed")
     }
 }