fix return usage segment

Author: brindy

Core/ReturnUserMeasurement.swift

@@ -20,8 +20,12 @@
 import Foundation
 import BrowserServicesKit
 
+/// This is only intended to be used during the install (first run after downloading from the app store).
 protocol ReturnUserMeasurement {
 
+    /// Based on the value in the keychain, so if you use this after the install process it will return true.
+    ///  If you really want to know if the user is "returning" then look at the variant in the `StatisticsStore`
+    ///  which will be set to `ru`.
     var isReturningUser: Bool { get }
     func installCompletedWithATB(_ atb: Atb)
     func updateStoredATB(_ atb: Atb)

Core/StatisticsLoader.swift

@@ -155,9 +155,10 @@ public class StatisticsLoader {
 
     private func processUsageSegmentation(atb: Atb?, activityType: UsageActivityType) {
         guard let installAtbValue = statisticsStore.atb else { return }
-        let installAtb = Atb(version: installAtbValue, updateVersion: nil)
-        let actualAtb = atb ?? installAtb
-        self.usageSegmentation.processATB(actualAtb, withInstallAtb: installAtb, andActivityType: activityType)
+        let installAtb = Atb(version: installAtbValue + (statisticsStore.variant ?? ""), updateVersion: nil)
+        let usageAtb = atb ?? installAtb
+
+        self.usageSegmentation.processATB(usageAtb, withInstallAtb: installAtb, andActivityType: activityType)
     }
 
     private func updateUsageSegmentationWithAtb(_ atb: Atb, activityType: UsageActivityType) {

DuckDuckGoTests/StatisticsLoaderTests.swift

@@ -42,6 +42,34 @@ class StatisticsLoaderTests: XCTestCase {
         super.tearDown()
     }
 
+    func testWhenAppRefreshHappensButNotInstalledAndReturningUser_ThenRetentionSegmentationNotified() {
+        mockStatisticsStore.variant = "ru"
+        mockStatisticsStore.atb = "v101-1"
+        
+        loadSuccessfulExiStub()
+
+        let testExpectation = expectation(description: "refresh complete")
+        testee.refreshAppRetentionAtb {
+            testExpectation.fulfill()
+        }
+        wait(for: [testExpectation], timeout: 5.0)
+        XCTAssertTrue(mockUsageSegmentation.atbs[0].installAtb.isReturningUser)
+    }
+
+    func testWhenReturnUser_ThenSegmentationIncludesCorrectVariant() {
+        mockStatisticsStore.variant = "ru"
+        mockStatisticsStore.atb = "v101-1"
+        mockStatisticsStore.searchRetentionAtb = "v101-2"
+        loadSuccessfulAtbStub()
+
+        let testExpectation = expectation(description: "refresh complete")
+        testee.refreshSearchRetentionAtb {
+            testExpectation.fulfill()
+        }
+        wait(for: [testExpectation], timeout: 5.0)
+        XCTAssertTrue(mockUsageSegmentation.atbs[0].installAtb.isReturningUser)
+    }
+
     func testWhenSearchRefreshHappensButNotInstalled_ThenRetentionSegmentationNotified() {
         loadSuccessfulExiStub()