From ccb9e4fdf2871308c4085e42e8b4e12c9906cea8 Mon Sep 17 00:00:00 2001
From: Chris Leonavicius <chris.leonavicius@amplitude.com>
Date: Wed, 2 Oct 2024 11:49:53 -0700
Subject: [PATCH] feat: add middleware removed callback

---
 Sources/Amplitude/AMPMiddleware.m        | 6 ++++++
 Sources/Amplitude/AMPMiddlewareRunner.h  | 3 +++
 Sources/Amplitude/AMPMiddlewareRunner.m  | 8 ++++++++
 Sources/Amplitude/Amplitude.m            | 1 +
 Sources/Amplitude/Public/AMPMiddleware.h | 2 ++
 Tests/AmplitudeTests.m                   | 6 ++++++
 6 files changed, 26 insertions(+)

diff --git a/Sources/Amplitude/AMPMiddleware.m b/Sources/Amplitude/AMPMiddleware.m
index 1b73fc2b..a1bd774d 100644
--- a/Sources/Amplitude/AMPMiddleware.m
+++ b/Sources/Amplitude/AMPMiddleware.m
@@ -91,4 +91,10 @@ - (void)amplitude:(Amplitude *)amplitude didOptOut:(BOOL)optOut {
     }
 }
 
+- (void)amplitudeDidRemoveMiddleware:(Amplitude *)amplitude {
+    if (self.didRemoveMiddleware) {
+        self.didRemoveMiddleware(amplitude);
+    }
+}
+
 @end
diff --git a/Sources/Amplitude/AMPMiddlewareRunner.h b/Sources/Amplitude/AMPMiddlewareRunner.h
index 51802363..28e60204 100644
--- a/Sources/Amplitude/AMPMiddlewareRunner.h
+++ b/Sources/Amplitude/AMPMiddlewareRunner.h
@@ -47,4 +47,7 @@
 - (void)dispatchAmplitude:(nonnull Amplitude *)amplitude didChangeUserId:(nonnull NSString *)userId;
 - (void)dispatchAmplitude:(nonnull Amplitude *)amplitude didOptOut:(BOOL)optOut;
 
+- (void)dispatchAmplitudeDidRemoveMiddleware:(nonnull Amplitude *)amplitude
+                                toMiddleware:(nonnull id<AMPMiddleware>)middleware;
+
 @end
diff --git a/Sources/Amplitude/AMPMiddlewareRunner.m b/Sources/Amplitude/AMPMiddlewareRunner.m
index 092fc146..126098eb 100644
--- a/Sources/Amplitude/AMPMiddlewareRunner.m
+++ b/Sources/Amplitude/AMPMiddlewareRunner.m
@@ -126,6 +126,14 @@ - (void)dispatchAmplitude:(Amplitude *)amplitude didOptOut:(BOOL)optOut {
     }
 }
 
+- (void)dispatchAmplitudeDidRemoveMiddleware:(Amplitude *)amplitude
+                                toMiddleware:(id<AMPMiddleware>)middleware {
+    if ([AMPMiddlewareRunner object:middleware
+                 respondsToSelector:@selector(amplitudeDidRemoveMiddleware:)]) {
+        [middleware amplitudeDidRemoveMiddleware:amplitude];
+    }
+}
+
 // AMPMiddleware never conformed to NSObject, which means we can't use the standard
 // [object respondsToSelector:] syntax to check for protocol conformance to optional methods.
 + (BOOL)object:(id)object respondsToSelector:(SEL)selector {
diff --git a/Sources/Amplitude/Amplitude.m b/Sources/Amplitude/Amplitude.m
index 2fe66ace..d08f940d 100644
--- a/Sources/Amplitude/Amplitude.m
+++ b/Sources/Amplitude/Amplitude.m
@@ -1668,6 +1668,7 @@ - (void)addEventMiddleware:(id<AMPMiddleware> _Nonnull)middleware {
 }
 
 - (void)removeEventMiddleware:(id<AMPMiddleware>)middleware {
+    [_middlewareRunner dispatchAmplitudeDidRemoveMiddleware:self toMiddleware:middleware];
     [_middlewareRunner remove:middleware];
 }
 
diff --git a/Sources/Amplitude/Public/AMPMiddleware.h b/Sources/Amplitude/Public/AMPMiddleware.h
index 598158c0..79a7fb09 100644
--- a/Sources/Amplitude/Public/AMPMiddleware.h
+++ b/Sources/Amplitude/Public/AMPMiddleware.h
@@ -54,6 +54,7 @@ typedef void (^AMPMiddlewareNext)(AMPMiddlewarePayload *_Nullable newPayload);
 - (void)amplitude:(nonnull Amplitude *)amplitude didChangeSessionId:(long long)sessionId;
 - (void)amplitude:(nonnull Amplitude *)amplitude didChangeUserId:(nonnull NSString *)userId;
 - (void)amplitude:(nonnull Amplitude *)amplitude didOptOut:(BOOL)optOut;
+- (void)amplitudeDidRemoveMiddleware:(nonnull Amplitude *)amplitude;
 
 @end
 
@@ -72,6 +73,7 @@ typedef void (^AMPMiddlewareBlock)(AMPMiddlewarePayload *_Nonnull payload, AMPMi
 @property (nonatomic, copy, nullable) void (^didChangeSessionId)(Amplitude * _Nonnull amplitude, long long sessionId);
 @property (nonatomic, copy, nullable) void (^didChangeUserId)(Amplitude * _Nonnull amplitude, NSString * _Nonnull userId);
 @property (nonatomic, copy, nullable) void (^didOptOut)(Amplitude * _Nonnull amplitude, BOOL optOut);
+@property (nonatomic, copy, nullable) void (^didRemoveMiddleware)(Amplitude * _Nonnull amplitude);
 
 - (instancetype _Nonnull)initWithBlock:(AMPMiddlewareBlock _Nonnull)block NS_DESIGNATED_INITIALIZER;
 
diff --git a/Tests/AmplitudeTests.m b/Tests/AmplitudeTests.m
index 40d1b858..0b0a2a04 100644
--- a/Tests/AmplitudeTests.m
+++ b/Tests/AmplitudeTests.m
@@ -1444,6 +1444,12 @@ - (void)testRemoveMiddleware {
         [receivedEventExpectation fulfill];
         next(payload);
     }];
+
+    const XCTestExpectation *didRemoveExpectation = [self expectationWithDescription:@"Removed middleware"];
+    middleware.didRemoveMiddleware = ^(Amplitude *amplitude) {
+        [didRemoveExpectation fulfill];
+    };
+
     [client addEventMiddleware:middleware];
 
     [client logEvent:@"test"];