From 03485578e53703720ac99b7a883bea9c38b1e56a Mon Sep 17 00:00:00 2001
From: Ariel Demarco <ariel.demarco@embrace.io>
Date: Tue, 18 Feb 2025 12:07:40 -0300
Subject: [PATCH] Fixed issue with swizzling

---
 ...sionDelegateProxyAsTaskDelegateTests.swift | 40 +++++++++++++------
 .../URLSessionCaptureServiceTests.swift       |  7 ++++
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/Tests/EmbraceCoreTests/Capture/Network/Proxy/URLSessionDelegateProxyAsTaskDelegateTests.swift b/Tests/EmbraceCoreTests/Capture/Network/Proxy/URLSessionDelegateProxyAsTaskDelegateTests.swift
index 13c57ee0..3b82f492 100644
--- a/Tests/EmbraceCoreTests/Capture/Network/Proxy/URLSessionDelegateProxyAsTaskDelegateTests.swift
+++ b/Tests/EmbraceCoreTests/Capture/Network/Proxy/URLSessionDelegateProxyAsTaskDelegateTests.swift
@@ -18,19 +18,6 @@ final class URLSessionDelegateProxyAsTaskDelegateTests: XCTestCase {
     static let timeoutQuick = 0.2
 
     // MARK: - Setup
-
-    override func tearDown() async throws {
-        do {
-            urlSessionCaptureService.swizzlers.forEach { swizzler in
-                try? swizzler.unswizzleClassMethod()
-                try? swizzler.unswizzleInstanceMethod()
-            }
-            try otherSwizzler?.unswizzleClassMethod()
-        } catch let exception {
-            print(exception)
-        }
-    }
-
     func givenCaptureServiceInstalled() {
         urlSessionCaptureService = URLSessionCaptureService(options: .init())
         openTelemetry = MockEmbraceOpenTelemetry()
@@ -59,6 +46,18 @@ final class URLSessionDelegateProxyAsTaskDelegateTests: XCTestCase {
         return url
     }
 
+    func unswizzleDefaultCaptureService() {
+        urlSessionCaptureService.swizzlers.forEach { swizzler in
+            try? swizzler.unswizzleClassMethod()
+            try? swizzler.unswizzleInstanceMethod()
+        }
+    }
+
+    func unswizzleOtherSwizzler() {
+        try? otherSwizzler?.unswizzleClassMethod()
+        try? otherSwizzler?.unswizzleInstanceMethod()
+    }
+
     // MARK: - Assertions
 
     /// Methods dealing with URLSessionTaskDelegate
@@ -83,6 +82,8 @@ final class URLSessionDelegateProxyAsTaskDelegateTests: XCTestCase {
             sessionDelegate.didReceiveDataExpectation,
             sessionDelegate.didCompleteWithErrorExpectation
         ], timeout: Self.timeoutQuick)
+
+        unswizzleDefaultCaptureService()
     }
 
     func givenSomebodyElseSwizzlesURLSessionInit() throws {
@@ -123,6 +124,9 @@ final class URLSessionDelegateProxyAsTaskDelegateTests: XCTestCase {
         XCTAssertTrue(try XCTUnwrap(otherSwizzler?.proxy?.didInvokeRespondsTo))
         XCTAssertTrue(try XCTUnwrap(otherSwizzler?.proxy?.didInvokeForwardingTarget))
         XCTAssertTrue(try XCTUnwrap(otherSwizzler?.proxy?.didForwardToTargetSuccessfully))
+
+        unswizzleDefaultCaptureService()
+        unswizzleOtherSwizzler()
     }
 
     @available(iOS 15.0, watchOS 8.0, *)
@@ -158,6 +162,9 @@ final class URLSessionDelegateProxyAsTaskDelegateTests: XCTestCase {
         XCTAssertTrue(try XCTUnwrap(otherSwizzler?.proxy?.didInvokeRespondsTo))
         XCTAssertTrue(try XCTUnwrap(otherSwizzler?.proxy?.didInvokeForwardingTarget))
         XCTAssertTrue(try XCTUnwrap(otherSwizzler?.proxy?.didForwardToTargetSuccessfully))
+
+        unswizzleOtherSwizzler()
+        unswizzleDefaultCaptureService()
     }
 
     @available(iOS 15.0, watchOS 8.0, *)
@@ -182,6 +189,7 @@ final class URLSessionDelegateProxyAsTaskDelegateTests: XCTestCase {
             taskDelegate.didReceiveDataExpectation,
             taskDelegate.didCompleteWithErrorExpectation
         ], timeout: Self.timeoutQuick)
+        unswizzleDefaultCaptureService()
     }
 
     @available(iOS 15.0, watchOS 8.0, *)
@@ -201,6 +209,8 @@ final class URLSessionDelegateProxyAsTaskDelegateTests: XCTestCase {
             sessionDelegate.didReceiveDataExpectation,
             sessionDelegate.didCompleteWithErrorExpectation
         ], timeout: Self.timeoutQuick)
+        
+        unswizzleDefaultCaptureService()
     }
 
     @available(iOS 15.0, watchOS 8.0, *)
@@ -218,6 +228,8 @@ final class URLSessionDelegateProxyAsTaskDelegateTests: XCTestCase {
 
         // DEV: async/await calls do not call `didCompleteWithError` method as response is handled inline
         XCTAssertFalse(sessionDelegate.didCallDidCompleteWithError)
+
+        unswizzleDefaultCaptureService()
     }
 
     @available(iOS 15.0, watchOS 8.0, *)
@@ -242,5 +254,7 @@ final class URLSessionDelegateProxyAsTaskDelegateTests: XCTestCase {
         XCTAssertFalse(sessionDelegate.didCallDidFinishCollecting)
         // DEV: async/await calls do not call `didCompleteWithError` method as response is handled inline
         XCTAssertFalse(sessionDelegate.didCallDidCompleteWithError)
+
+        unswizzleDefaultCaptureService()
     }
 }
diff --git a/Tests/EmbraceCoreTests/Capture/Network/URLSessionCaptureServiceTests.swift b/Tests/EmbraceCoreTests/Capture/Network/URLSessionCaptureServiceTests.swift
index a898a198..d1e453b5 100644
--- a/Tests/EmbraceCoreTests/Capture/Network/URLSessionCaptureServiceTests.swift
+++ b/Tests/EmbraceCoreTests/Capture/Network/URLSessionCaptureServiceTests.swift
@@ -21,6 +21,13 @@ class URLSessionCaptureServiceTests: XCTestCase {
         givenURLSessionSwizzlerProvider()
     }
 
+    override func tearDown() {
+        sut.swizzlers.forEach {
+            try? $0.unswizzleClassMethod()
+            try? $0.unswizzleInstanceMethod()
+        }
+    }
+
     func test_onInit_collectorIsUninstalled() {
         whenInitializingURLSessionCaptureService()
         thenCaptureServiceStatus(is: .uninstalled)