feat: Add more functionality to middleware (#502)

amplitude · Aug 24, 2024 · 5748e96 · 5748e96
1 parent b5bc44c
commit 5748e96
Show file tree

Hide file tree

Showing 7 changed files with 221 additions and 2 deletions.
diff --git a/Sources/Amplitude/AMPMiddleware.m b/Sources/Amplitude/AMPMiddleware.m
@@ -67,4 +67,28 @@ - (void)amplitude:(Amplitude *)amplitude didUploadEventsManually:(BOOL)manually
     }
 }
 
+- (void)amplitude:(Amplitude *)amplitude didChangeDeviceId:(NSString *)deviceId {
+    if (self.didChangeDeviceId) {
+        self.didChangeDeviceId(amplitude, deviceId);
+    }
+}
+
+- (void)amplitude:(Amplitude *)amplitude didChangeSessionId:(long long)sessionId {
+    if (self.didChangeSessionId) {
+        self.didChangeSessionId(amplitude, sessionId);
+    }
+}
+
+- (void)amplitude:(Amplitude *)amplitude didChangeUserId:(NSString *)userId {
+    if (self.didChangeUserId) {
+        self.didChangeUserId(amplitude, userId);
+    }
+}
+
+- (void)amplitude:(Amplitude *)amplitude didOptOut:(BOOL)optOut {
+    if (self.didOptOut) {
+        self.didOptOut(amplitude, optOut);
+    }
+}
+
 @end
diff --git a/Sources/Amplitude/AMPMiddlewareRunner.h b/Sources/Amplitude/AMPMiddlewareRunner.h
@@ -33,13 +33,18 @@
 
 - (void) add:(id<AMPMiddleware> _Nonnull)middleware;
 
+- (void)remove:(nonnull id<AMPMiddleware>)middleware;
+
 - (void) run:(AMPMiddlewarePayload *_Nonnull)payload next:(AMPMiddlewareNext _Nonnull)next;
 
 - (void)dispatchAmplitudeInitialized:(nonnull Amplitude *)amplitude;
 - (void)dispatchAmplitudeInitialized:(nonnull Amplitude *)amplitude
                         toMiddleware:(nonnull id<AMPMiddleware>)middleware;
 
-- (void)dispatchAmplitude:(nonnull Amplitude *)amplitude
-  didUploadEventsManually:(BOOL)isManualUpload;
+- (void)dispatchAmplitude:(nonnull Amplitude *)amplitude didUploadEventsManually:(BOOL)isManualUpload;
+- (void)dispatchAmplitude:(nonnull Amplitude *)amplitude didChangeDeviceId:(nonnull NSString *)deviceId;
+- (void)dispatchAmplitude:(nonnull Amplitude *)amplitude didChangeSessionId:(long long)sessionId;
+- (void)dispatchAmplitude:(nonnull Amplitude *)amplitude didChangeUserId:(nonnull NSString *)userId;
+- (void)dispatchAmplitude:(nonnull Amplitude *)amplitude didOptOut:(BOOL)optOut;
 
 @end
diff --git a/Sources/Amplitude/AMPMiddlewareRunner.m b/Sources/Amplitude/AMPMiddlewareRunner.m
@@ -43,6 +43,10 @@ - (void) add:(id<AMPMiddleware> _Nonnull)middleware {
     [self.middlewares addObject:middleware];
 }
 
+- (void)remove:(id<AMPMiddleware>)middleware {
+    [self.middlewares removeObject:middleware];
+}
+
 - (void) run:(AMPMiddlewarePayload *_Nonnull)payload next:(AMPMiddlewareNext _Nonnull)next {
     [self runMiddlewares:self.middlewares payload:payload callback:next];
 }
@@ -86,6 +90,42 @@ - (void)dispatchAmplitude:(Amplitude *)amplitude didUploadEventsManually:(BOOL)i
     }
 }
 
+- (void)dispatchAmplitude:(Amplitude *)amplitude didChangeDeviceId:(NSString *)deviceId {
+    for (id<AMPMiddleware> middleware in self.middlewares) {
+        if ([AMPMiddlewareRunner object:middleware
+                     respondsToSelector:@selector(amplitude:didChangeDeviceId:)]) {
+            [middleware amplitude:amplitude didChangeDeviceId:deviceId];
+        }
+    }
+}
+
+- (void)dispatchAmplitude:(Amplitude *)amplitude didChangeSessionId:(long long)sessionId {
+    for (id<AMPMiddleware> middleware in self.middlewares) {
+        if ([AMPMiddlewareRunner object:middleware
+                     respondsToSelector:@selector(amplitude:didChangeSessionId:)]) {
+            [middleware amplitude:amplitude didChangeSessionId:sessionId];
+        }
+    }
+}
+
+- (void)dispatchAmplitude:(Amplitude *)amplitude didChangeUserId:(NSString *)userId {
+    for (id<AMPMiddleware> middleware in self.middlewares) {
+        if ([AMPMiddlewareRunner object:middleware
+                     respondsToSelector:@selector(amplitude:didChangeUserId:)]) {
+            [middleware amplitude:amplitude didChangeUserId:userId];
+        }
+    }
+}
+
+- (void)dispatchAmplitude:(Amplitude *)amplitude didOptOut:(BOOL)optOut {
+    for (id<AMPMiddleware> middleware in self.middlewares) {
+        if ([AMPMiddlewareRunner object:middleware
+                     respondsToSelector:@selector(amplitude:didOptOut:)]) {
+            [middleware amplitude:amplitude didOptOut:optOut];
+        }
+    }
+}
+
 // 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
@@ -1384,6 +1384,7 @@ - (BOOL)isWithinMinTimeBetweenSessions:(NSNumber *)timestamp {
 - (void)setSessionId:(long long)timestamp {
     _sessionId = timestamp;
     [self setPreviousSessionId:_sessionId];
+    [_middlewareRunner dispatchAmplitude:self didChangeSessionId:timestamp];
 }
 
 /**
@@ -1548,6 +1549,7 @@ - (void)setUserId:(NSString *)userId startNewSession:(BOOL)startNewSession {
         // between Analytics and Experiment SDKs.
         id<IdentityStoreEditor> identityStoreEditor = [[[AnalyticsConnector getInstance:self.instanceName] identityStore] editIdentity];
         [[identityStoreEditor setUserId:self.userId] commit];
+        [self->_middlewareRunner dispatchAmplitude:self didChangeUserId:self.userId];
 
         if (startNewSession) {
             NSNumber *timestamp = [NSNumber numberWithLongLong:[[self currentTime] timeIntervalSince1970] * 1000];
@@ -1564,6 +1566,7 @@ - (void)setOptOut:(BOOL)enabled {
     [self runOnBackgroundQueue:^{
         NSNumber *value = [NSNumber numberWithBool:enabled];
         (void) [self.dbHelper insertOrReplaceKeyLongValue:OPT_OUT value:value];
+        [self->_middlewareRunner dispatchAmplitude:self didOptOut:enabled];
     }];
 }
 
@@ -1620,6 +1623,7 @@ - (void)setDeviceId:(NSString *)deviceId {
         // between Analytics and Experiment SDKs.
         id<IdentityStoreEditor> identityStoreEditor = [[[AnalyticsConnector getInstance:self.instanceName] identityStore] editIdentity];
         [[identityStoreEditor setDeviceId:self.deviceId] commit];
+        [self->_middlewareRunner dispatchAmplitude:self didChangeDeviceId:self.deviceId];
     }];
 }
 
@@ -1663,6 +1667,10 @@ - (void)addEventMiddleware:(id<AMPMiddleware> _Nonnull)middleware {
     }
 }
 
+- (void)removeEventMiddleware:(id<AMPMiddleware>)middleware {
+    [_middlewareRunner remove:middleware];
+}
+
 /**
  * The amount of time after an identify is logged that identify events will be batched before being uploaded to the server.
  * The default is 30 seconds.

diff --git a/Sources/Amplitude/Public/AMPMiddleware.h b/Sources/Amplitude/Public/AMPMiddleware.h
@@ -50,6 +50,10 @@ typedef void (^AMPMiddlewareNext)(AMPMiddlewarePayload *_Nullable newPayload);
 
 - (void)amplitudeDidFinishInitializing:(nonnull Amplitude *)amplitude;
 - (void)amplitude:(nonnull Amplitude *)amplitude didUploadEventsManually:(BOOL)manually;
+- (void)amplitude:(nonnull Amplitude *)amplitude didChangeDeviceId:(nonnull NSString *)deviceId;
+- (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;
 
 @end
 
@@ -64,6 +68,10 @@ typedef void (^AMPMiddlewareBlock)(AMPMiddlewarePayload *_Nonnull payload, AMPMi
 
 @property (nonatomic, copy, nullable) void (^didFinishInitializing)(Amplitude * _Nonnull amplitude);
 @property (nonatomic, copy, nullable) void (^didUploadEventsManually)(Amplitude * _Nonnull amplitude, BOOL isManualUpload);
+@property (nonatomic, copy, nullable) void (^didChangeDeviceId)(Amplitude * _Nonnull amplitude, NSString * _Nonnull deviceId);
+@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);
 
 - (instancetype _Nonnull)initWithBlock:(AMPMiddlewareBlock _Nonnull)block NS_DESIGNATED_INITIALIZER;
 

diff --git a/Sources/Amplitude/Public/Amplitude.h b/Sources/Amplitude/Public/Amplitude.h
@@ -698,6 +698,11 @@ typedef void (^AMPInitCompletionBlock)(void);
  */
 - (void)addEventMiddleware:(id<AMPMiddleware> _Nonnull)middleware;
 
+/**
+ * Removes an existing middleware function to run on each logEvent() call prior to sending to Amplitude.
+ */
+- (void)removeEventMiddleware:(id<AMPMiddleware> _Nonnull)middleware;
+
 /**
  * The amount of time after an identify is logged that identify events will be batched before being uploaded to the server.
  * The default is 30 seconds.

diff --git a/Tests/AmplitudeTests.m b/Tests/AmplitudeTests.m
@@ -1322,6 +1322,107 @@ - (void)testMiddlewareAutomaticFlush {
     }];
 }
 
+- (void)testMiddlewareDeviceIdChanged {
+    const Amplitude *client = [Amplitude instanceWithName:@"middleware_device_id_changed"];
+    [client setDeviceId:@"old_id"];
+    [client initializeApiKey:@"aaa"];
+    [client flushQueue];
+
+    NSString *expectedDeviceId = @"updated_device_id";
+
+    const XCTestExpectation *deviceIdChangedExpectation = [self expectationWithDescription:@"changed device id"];
+    const AMPBlockMiddleware *deviceIdChangedMiddleware = [[AMPBlockMiddleware alloc] init];
+    deviceIdChangedMiddleware.didChangeDeviceId = ^(Amplitude *amplitude, NSString *deviceId) {
+        XCTAssertEqual(amplitude.deviceId, expectedDeviceId);
+        XCTAssertEqual(deviceId, expectedDeviceId);
+        [deviceIdChangedExpectation fulfill];
+    };
+    [client addEventMiddleware:deviceIdChangedMiddleware];
+
+    [client setDeviceId:expectedDeviceId];
+
+    [self waitForExpectationsWithTimeout:10.0 handler:^(NSError * _Nullable error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
+- (void)testMiddlewareSessionIdChanged {
+    const Amplitude *client = [Amplitude instanceWithName:@"middleware_session_id_changed"];
+    [client setSessionId:1];
+    [client initializeApiKey:@"aaa"];
+    [client flushQueue];
+
+    long long expectedSessionId = [NSDate date].timeIntervalSince1970 * 1000;
+
+    const XCTestExpectation *sessionIdChangedExpectation = [self expectationWithDescription:@"changed session id"];
+    const AMPBlockMiddleware *userIdChangedMiddleware = [[AMPBlockMiddleware alloc] init];
+    userIdChangedMiddleware.didChangeSessionId = ^(Amplitude *amplitude, long long sessionId) {
+        XCTAssertEqual(amplitude.sessionId, expectedSessionId);
+        XCTAssertEqual(sessionId, expectedSessionId);
+        [sessionIdChangedExpectation fulfill];
+    };
+    [client addEventMiddleware:userIdChangedMiddleware];
+
+    [client setSessionId:expectedSessionId];
+
+    [self waitForExpectationsWithTimeout:10.0 handler:^(NSError * _Nullable error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
+- (void)testMiddlewareUserIdChanged {
+    const Amplitude *client = [Amplitude instanceWithName:@"middleware_user_id_changed"];
+    [client initializeApiKey:@"aaa" userId:@"old_user_id"];
+    [client flushQueue];
+
+    NSString *expectedUserId = @"updated_user_id";
+
+    const XCTestExpectation *userIdChangedExpectation = [self expectationWithDescription:@"changed user id"];
+    const AMPBlockMiddleware *userIdChangedMiddleware = [[AMPBlockMiddleware alloc] init];
+    userIdChangedMiddleware.didChangeUserId = ^(Amplitude *amplitude, NSString *userId) {
+        XCTAssertEqual(amplitude.userId, expectedUserId);
+        XCTAssertEqual(userId, expectedUserId);
+        [userIdChangedExpectation fulfill];
+    };
+    [client addEventMiddleware:userIdChangedMiddleware];
+
+    [client setUserId:expectedUserId];
+
+    [self waitForExpectationsWithTimeout:10.0 handler:^(NSError * _Nullable error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
+- (void)testMiddlewareOptOutChanged {
+    const Amplitude *client = [Amplitude instanceWithName:@"middleware_opt_out_changed"];
+    client.optOut = NO;
+    [client initializeApiKey:@"aaa"];
+    [client flushQueue];
+
+    const XCTestExpectation *optOutChangedExpectation = [self expectationWithDescription:@"changed user id"];
+    const AMPBlockMiddleware *optOutChangedMiddleware = [[AMPBlockMiddleware alloc] init];
+    optOutChangedMiddleware.didOptOut = ^(Amplitude *amplitude, BOOL optOut) {
+        XCTAssertTrue(optOut);
+        XCTAssertTrue(amplitude.optOut);
+        [optOutChangedExpectation fulfill];
+    };
+    [client addEventMiddleware:optOutChangedMiddleware];
+
+    client.optOut = YES;
+
+    [self waitForExpectationsWithTimeout:10.0 handler:^(NSError * _Nullable error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
 - (void)testSwallowMiddleware {
     AMPBlockMiddleware *swallowMiddleware = [[AMPBlockMiddleware alloc] initWithBlock: ^(AMPMiddlewarePayload * _Nonnull payload, AMPMiddlewareNext _Nonnull next) {
     }];
@@ -1334,6 +1435,34 @@ - (void)testSwallowMiddleware {
     XCTAssertNil([client getLastEventFromInstanceName:@"middleware_swallow"]);
 }
 
+- (void)testRemoveMiddleware {
+    Amplitude *client = [Amplitude instanceWithName:@"remove_middleware"];
+    [client initializeApiKey:@"remove_middleware_api_key"];
+
+    const XCTestExpectation *receivedEventExpectation = [self expectationWithDescription:@"Received event"];
+    const AMPBlockMiddleware *middleware = [[AMPBlockMiddleware alloc] initWithBlock:^(AMPMiddlewarePayload * _Nonnull payload, AMPMiddlewareNext  _Nonnull next) {
+        [receivedEventExpectation fulfill];
+        next(payload);
+    }];
+    [client addEventMiddleware:middleware];
+
+    [client logEvent:@"test"];
+
+    [client flushQueue];
+
+    [client removeEventMiddleware:middleware];
+
+    [client logEvent:@"test2" ];
+
+    [client flushQueue];
+
+    [self waitForExpectationsWithTimeout:10.0 handler:^(NSError * _Nullable error) {
+        if (error) {
+            XCTFail(@"Timeout");
+        }
+    }];
+}
+
 -(void)testLogEventWithUserProperties {
     NSString *instanceName = @"eventWithUserProperties";
     Amplitude *client = [Amplitude instanceWithName:instanceName];