From 8b5801c82b5f1866b787adc36fcd6144feff745e Mon Sep 17 00:00:00 2001
From: FitseTLT <engfitsum@gmail.com>
Date: Wed, 13 Aug 2025 20:07:20 +0300
Subject: [PATCH] added the new policy data on snapshot when changing workspace

---
 src/libs/actions/Report.ts  | 15 +++++++++++++++
 tests/actions/ReportTest.ts | 29 ++++++++++++++++++++++++++++-
 2 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/libs/actions/Report.ts b/src/libs/actions/Report.ts
index f784be4a7ac2..fc1ddd3043b6 100644
--- a/src/libs/actions/Report.ts
+++ b/src/libs/actions/Report.ts
@@ -157,6 +157,7 @@ import {
     isValidReportIDFromPath,
     prepareOnboardingOnyxData,
 } from '@libs/ReportUtils';
+import {getCurrentSearchQueryJSON} from '@libs/SearchQueryUtils';
 import shouldSkipDeepLinkNavigation from '@libs/shouldSkipDeepLinkNavigation';
 import playSound, {SOUNDS} from '@libs/Sound';
 import {isOnHold} from '@libs/TransactionUtils';
@@ -5756,6 +5757,20 @@ function buildOptimisticChangePolicyData(report: Report, policyID: string, repor
         },
     });
 
+    const currentSearchQueryJSON = getCurrentSearchQueryJSON();
+    const policy = allPolicies?.[`${ONYXKEYS.COLLECTION.POLICY}${policyID}`];
+
+    // Search data might not have the new policy data so we should add it optimistically.
+    if (policy && currentSearchQueryJSON) {
+        optimisticData.push({
+            onyxMethod: Onyx.METHOD.MERGE,
+            key: `${ONYXKEYS.COLLECTION.SNAPSHOT}${currentSearchQueryJSON.hash}` as const,
+            value: {
+                data: {[`${ONYXKEYS.COLLECTION.POLICY}${policy.id}`]: policy},
+            },
+        });
+    }
+
     // 5. Make sure the expense report is not archived
     optimisticData.push({
         onyxMethod: Onyx.METHOD.MERGE,
diff --git a/tests/actions/ReportTest.ts b/tests/actions/ReportTest.ts
index 8c914d889df6..a67fe960a66b 100644
--- a/tests/actions/ReportTest.ts
+++ b/tests/actions/ReportTest.ts
@@ -44,6 +44,17 @@ jest.mock('@libs/ReportUtils', () => {
     };
 });
 
+const currentHash = 12345;
+jest.mock('@src/libs/SearchQueryUtils', () => ({
+    getCurrentSearchQueryJSON: jest.fn().mockImplementation(() => ({
+        hash: currentHash,
+        query: 'test',
+        type: 'expense',
+        status: '',
+        flatFilters: [],
+    })),
+}));
+
 const UTC = 'UTC';
 jest.mock('@src/libs/actions/Report', () => {
     const originalModule = jest.requireActual<Report>('@src/libs/actions/Report');
@@ -1835,8 +1846,11 @@ describe('actions/Report', () => {
                 private_isArchived: DateUtils.getDBTime(),
             });
 
+            const newPolicy = createRandomPolicy(2);
+            await Onyx.merge(`${ONYXKEYS.COLLECTION.POLICY}${newPolicy.id}`, newPolicy);
+
             // When moving to another workspace
-            Report.changeReportPolicy(expenseReport, '2');
+            Report.changeReportPolicy(expenseReport, newPolicy.id);
             await waitForBatchedUpdates();
 
             // Then the expense report should not be archived anymore
@@ -1850,6 +1864,19 @@ describe('actions/Report', () => {
                 });
             });
             expect(isArchived).toBe(false);
+
+            const snapshotData = await new Promise<OnyxEntry<OnyxTypes.SearchResults>>((resolve) => {
+                const connection = Onyx.connect({
+                    key: `${ONYXKEYS.COLLECTION.SNAPSHOT}${currentHash}`,
+                    callback: (val) => {
+                        resolve(val);
+                        Onyx.disconnect(connection);
+                    },
+                });
+            });
+
+            // Then the new policy data should also be populated on the current search snapshot.
+            expect(snapshotData?.data?.[`${ONYXKEYS.COLLECTION.POLICY}${newPolicy.id}`]).toBeDefined();
         });
     });