Fix issue where user-supplied enrichments were lost during the startu…

…p phase. (#367)
segmentio · Oct 8, 2024 · 483236e · 483236e
1 parent 41df310
commit 483236e
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 10 deletions.
diff --git a/Sources/Segment/Analytics.swift b/Sources/Segment/Analytics.swift
@@ -95,9 +95,9 @@ public class Analytics {
 
     internal func process<E: RawEvent>(incomingEvent: E, enrichments: [EnrichmentClosure]? = nil) {
         guard enabled == true else { return }
-        let event = incomingEvent.applyRawEventData(store: store)
+        let event = incomingEvent.applyRawEventData(store: store, enrichments: enrichments)
 
-        _ = timeline.process(incomingEvent: event, enrichments: enrichments)
+        _ = timeline.process(incomingEvent: event)
 
         /*let flushPolicies = configuration.values.flushPolicies
         for policy in flushPolicies {

diff --git a/Sources/Segment/Timeline.swift b/Sources/Segment/Timeline.swift
@@ -25,13 +25,13 @@ public class Timeline {
     }
 
     @discardableResult
-    internal func process<E: RawEvent>(incomingEvent: E, enrichments: [EnrichmentClosure]? = nil) -> E? {
+    internal func process<E: RawEvent>(incomingEvent: E) -> E? {
         // apply .before and .enrichment types first ...
         let beforeResult = applyPlugins(type: .before, event: incomingEvent)
         // .enrichment here is akin to source middleware in the old analytics-ios.
         var enrichmentResult = applyPlugins(type: .enrichment, event: beforeResult)
 
-        if let enrichments {
+        if let enrichments = enrichmentResult?.enrichments {
             for closure in enrichments {
                 if let result = closure(enrichmentResult) as? E {
                     enrichmentResult = result

diff --git a/Sources/Segment/Types.swift b/Sources/Segment/Types.swift
@@ -19,6 +19,7 @@ public struct DestinationMetadata: Codable {
 // MARK: - Event Types
 
 public protocol RawEvent: Codable {
+    var enrichments: [EnrichmentClosure]? { get set }
     var type: String? { get set }
     var anonymousId: String? { get set }
     var messageId: String? { get set }
@@ -32,6 +33,7 @@ public protocol RawEvent: Codable {
 }
 
 public struct TrackEvent: RawEvent {
+    @Noncodable public var enrichments: [EnrichmentClosure]? = nil
     public var type: String? = "track"
     public var anonymousId: String? = nil
     public var messageId: String? = nil
@@ -57,6 +59,7 @@ public struct TrackEvent: RawEvent {
 }
 
 public struct IdentifyEvent: RawEvent {
+    @Noncodable public var enrichments: [EnrichmentClosure]? = nil
     public var type: String? = "identify"
     public var anonymousId: String? = nil
     public var messageId: String? = nil
@@ -82,6 +85,7 @@ public struct IdentifyEvent: RawEvent {
 }
 
 public struct ScreenEvent: RawEvent {
+    @Noncodable public var enrichments: [EnrichmentClosure]? = nil
     public var type: String? = "screen"
     public var anonymousId: String? = nil
     public var messageId: String? = nil
@@ -109,6 +113,7 @@ public struct ScreenEvent: RawEvent {
 }
 
 public struct GroupEvent: RawEvent {
+    @Noncodable public var enrichments: [EnrichmentClosure]? = nil
     public var type: String? = "group"
     public var anonymousId: String? = nil
     public var messageId: String? = nil
@@ -134,6 +139,7 @@ public struct GroupEvent: RawEvent {
 }
 
 public struct AliasEvent: RawEvent {
+    @Noncodable public var enrichments: [EnrichmentClosure]? = nil
     public var type: String? = "alias"
     public var anonymousId: String? = nil
     public var messageId: String? = nil
@@ -289,11 +295,12 @@ extension RawEvent {
         }
     }
 
-    internal func applyRawEventData(store: Store) -> Self {
+    internal func applyRawEventData(store: Store, enrichments: [EnrichmentClosure]?) -> Self {
         var result: Self = self
 
         guard let userInfo: UserInfo = store.currentState() else { return self }
 
+        result.enrichments = enrichments
         result.anonymousId = userInfo.anonymousId
         result.userId = userInfo.userId
         result.messageId = UUID().uuidString

diff --git a/Sources/Segment/Utilities/Noncodable.swift b/Sources/Segment/Utilities/Noncodable.swift
@@ -8,7 +8,7 @@
 import Foundation
 
 @propertyWrapper
-internal struct Noncodable<T>: Codable {
+public struct Noncodable<T>: Codable {
     public var wrappedValue: T?
     public init(wrappedValue: T?) {
         self.wrappedValue = wrappedValue

diff --git a/Tests/Segment-Tests/Analytics_Tests.swift b/Tests/Segment-Tests/Analytics_Tests.swift
@@ -1017,18 +1017,23 @@ final class Analytics_Tests: XCTestCase {
         let analytics = Analytics(configuration: Configuration(writeKey: "test"))
         let outputReader = OutputReaderPlugin()
         analytics.add(plugin: outputReader)
-
-        waitUntilStarted(analytics: analytics)
 
         let addEventOrigin: EnrichmentClosure = { event in
             return Context.insertOrigin(event: event, data: [
                 "type": "mobile"
             ])
         }
+
+        analytics.track(name: "enrichment check pre startup", enrichments: [addEventOrigin])
+
+        waitUntilStarted(analytics: analytics)
+
+        let trackEvent1: TrackEvent? = outputReader.lastEvent as? TrackEvent
+        XCTAssertEqual(trackEvent1?.context?.value(forKeyPath: "__eventOrigin.type"), "mobile")
 
         analytics.track(name: "enrichment check", enrichments: [addEventOrigin])
 
-        let trackEvent: TrackEvent? = outputReader.lastEvent as? TrackEvent
-        XCTAssertEqual(trackEvent?.context?.value(forKeyPath: "__eventOrigin.type"), "mobile")
+        let trackEvent2: TrackEvent? = outputReader.lastEvent as? TrackEvent
+        XCTAssertEqual(trackEvent2?.context?.value(forKeyPath: "__eventOrigin.type"), "mobile")
     }
 }