feat: add middleware removed callback (#506)

amplitude · Oct 2, 2024 · b1b001a · b1b001a
1 parent 681c498
commit b1b001a
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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"];