From e6e47a2eebefbc5c64d552c611c2fc637454c93b Mon Sep 17 00:00:00 2001
From: Toomas Vahter <toomas.vahter@getstream.io>
Date: Tue, 10 Dec 2024 11:56:12 +0200
Subject: [PATCH] End background task as soon as begin task is called

---
 CHANGELOG.md                                  |  2 +
 .../BackgroundTaskScheduler.swift             |  1 +
 .../BackgroundTaskScheduler_Tests.swift       | 41 ++++++++++++-------
 3 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index dbac824faf6..708cc59d3c8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
   - Add `ChannelListSortingKey.pinnedAt`
   - Add `ChatChannel.membership.pinnedAt`
   - Add `ChatChannel.isPinned`
+### 🐞 Fixed
+- Fix a rare issue of not ending background tasks [#3528](https://github.com/GetStream/stream-chat-swift/pull/3528)
 
 # [4.68.0](https://github.com/GetStream/stream-chat-swift/releases/tag/4.68.0)
 _December 03, 2024_
diff --git a/Sources/StreamChat/WebSocketClient/BackgroundTaskScheduler.swift b/Sources/StreamChat/WebSocketClient/BackgroundTaskScheduler.swift
index 23fa1c68d2d..a22e6954725 100644
--- a/Sources/StreamChat/WebSocketClient/BackgroundTaskScheduler.swift
+++ b/Sources/StreamChat/WebSocketClient/BackgroundTaskScheduler.swift
@@ -51,6 +51,7 @@ class IOSBackgroundTaskScheduler: BackgroundTaskScheduler {
     }
 
     func beginTask(expirationHandler: (() -> Void)?) -> Bool {
+        endTask()
         activeBackgroundTask = app?.beginBackgroundTask { [weak self] in
             expirationHandler?()
             self?.endTask()
diff --git a/Tests/StreamChatTests/WebSocketClient/BackgroundTaskScheduler_Tests.swift b/Tests/StreamChatTests/WebSocketClient/BackgroundTaskScheduler_Tests.swift
index 1b0c9b7fb18..62e70c4000e 100644
--- a/Tests/StreamChatTests/WebSocketClient/BackgroundTaskScheduler_Tests.swift
+++ b/Tests/StreamChatTests/WebSocketClient/BackgroundTaskScheduler_Tests.swift
@@ -3,6 +3,7 @@
 //
 
 @testable import StreamChat
+import StreamChatTestTools
 import XCTest
 
 #if os(iOS)
@@ -44,22 +45,9 @@ final class IOSBackgroundTaskScheduler_Tests: XCTestCase {
     }
 
     func test_whenSchedulerIsDeallocated_backgroundTaskIsEnded() {
-        // Create mock scheduler type and catch `endTask` invokation
-        class MockScheduler: IOSBackgroundTaskScheduler {
-            let endTaskClosure: () -> Void
-
-            init(endTaskClosure: @escaping () -> Void) {
-                self.endTaskClosure = endTaskClosure
-            }
-
-            override func endTask() {
-                endTaskClosure()
-            }
-        }
-
         // Create mock scheduler and catch `endTask`
         var endTaskCalled = false
-        var scheduler: MockScheduler? = MockScheduler {
+        var scheduler: IOSBackgroundTaskSchedulerMock? = IOSBackgroundTaskSchedulerMock {
             endTaskCalled = true
         }
 
@@ -75,5 +63,30 @@ final class IOSBackgroundTaskScheduler_Tests: XCTestCase {
         // Simulate access to scheduler to eliminate the warning
         _ = scheduler
     }
+    
+    func test_callingBeginMultipleTimes_allTheBackgroundTasksAreEnded() {
+        var endTaskCallCount = 0
+        let scheduler = IOSBackgroundTaskSchedulerMock {
+            endTaskCallCount += 1
+        }
+        _ = scheduler.beginTask(expirationHandler: nil)
+        _ = scheduler.beginTask(expirationHandler: nil)
+        _ = scheduler.beginTask(expirationHandler: nil)
+        XCTAssertEqual(3, endTaskCallCount)
+    }
+    
+    // MARK: - Mocks
+    
+    class IOSBackgroundTaskSchedulerMock: IOSBackgroundTaskScheduler {
+        let endTaskClosure: () -> Void
+
+        init(endTaskClosure: @escaping () -> Void) {
+            self.endTaskClosure = endTaskClosure
+        }
+
+        override func endTask() {
+            endTaskClosure()
+        }
+    }
 }
 #endif