From 75b429d143d89e354cfa907d4b627621844564d4 Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Tue, 21 Mar 2023 12:13:20 -0700
Subject: [PATCH] cherry-pick(#21844): chore: update test locations when
 merging

---
 .../src/isomorphic/teleReceiver.ts            |  1 +
 packages/trace-viewer/src/ui/sourceTab.tsx    |  2 +-
 packages/trace-viewer/src/ui/watchMode.tsx    |  4 ++
 .../ui-mode-test-update.spec.ts               | 59 +++++++++++++++++++
 4 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/packages/playwright-test/src/isomorphic/teleReceiver.ts b/packages/playwright-test/src/isomorphic/teleReceiver.ts
index ba3a92df30..4402bfb8eb 100644
--- a/packages/playwright-test/src/isomorphic/teleReceiver.ts
+++ b/packages/playwright-test/src/isomorphic/teleReceiver.ts
@@ -301,6 +301,7 @@ export class TeleReporterReceiver {
     test.id = payload.testId;
     test.expectedStatus = payload.expectedStatus;
     test.timeout = payload.timeout;
+    test.location = payload.location;
     test.annotations = payload.annotations;
     test.retries = payload.retries;
     return test;
diff --git a/packages/trace-viewer/src/ui/sourceTab.tsx b/packages/trace-viewer/src/ui/sourceTab.tsx
index 3951ddde91..f0c7fb8600 100644
--- a/packages/trace-viewer/src/ui/sourceTab.tsx
+++ b/packages/trace-viewer/src/ui/sourceTab.tsx
@@ -59,7 +59,7 @@ export const SourceTab: React.FunctionComponent<{
     const highlight: SourceHighlight[] = source.errors.map(e => ({ type: 'error', line: e.location.line, message: e.error!.message }));
     highlight.push({ line: targetLine, type: 'running' });
 
-    if (source.content === undefined) {
+    if (source.content === undefined || fallbackLocation) {
       const sha1 = await calculateSha1(location.file);
       try {
         let response = await fetch(`sha1/src@${sha1}.txt`);
diff --git a/packages/trace-viewer/src/ui/watchMode.tsx b/packages/trace-viewer/src/ui/watchMode.tsx
index 6473e1fae7..cebd052e30 100644
--- a/packages/trace-viewer/src/ui/watchMode.tsx
+++ b/packages/trace-viewer/src/ui/watchMode.tsx
@@ -614,6 +614,10 @@ const sendMessage = async (method: string, params: any) => {
 };
 
 const sendMessageNoReply = (method: string, params?: any) => {
+  if ((window as any)._overrideProtocolForTest) {
+    (window as any)._overrideProtocolForTest({ method, params }).catch(() => {});
+    return;
+  }
   sendMessage(method, params).catch((e: Error) => {
     // eslint-disable-next-line no-console
     console.error(e);
diff --git a/tests/playwright-test/ui-mode-test-update.spec.ts b/tests/playwright-test/ui-mode-test-update.spec.ts
index 8d64a3b50c..3e3a1a53bc 100644
--- a/tests/playwright-test/ui-mode-test-update.spec.ts
+++ b/tests/playwright-test/ui-mode-test-update.spec.ts
@@ -168,3 +168,62 @@ test('should pick new / deleted nested tests', async ({ runUITest, writeFiles, d
           ◯ inner fails
   `);
 });
+
+test('should update test locations', async ({ runUITest, writeFiles, deleteFile }) => {
+  const page = await runUITest({
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('passes', () => {});
+    `,
+  });
+
+  await expect.poll(dumpTestTree(page), { timeout: 15000 }).toContain(`
+    ▼ ◯ a.test.ts
+        ◯ passes
+  `);
+
+  const messages: any = [];
+  await page.exposeBinding('_overrideProtocolForTest', (_, data) => messages.push(data));
+
+  const passesItemLocator = page.getByRole('listitem').filter({ hasText: 'passes' });
+  await passesItemLocator.hover();
+  await passesItemLocator.getByTitle('Open in VS Code').click();
+
+  expect(messages).toEqual([{
+    method: 'open',
+    params: {
+      location: expect.stringContaining('a.test.ts:3'),
+    },
+  }]);
+
+  await writeFiles({
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('new-test', () => {});
+
+      test('passes', () => {});
+    `
+  });
+
+  await expect.poll(dumpTestTree(page), { timeout: 15000 }).toContain(`
+    ▼ ◯ a.test.ts
+        ◯ new-test
+        ◯ passes <=
+  `);
+
+  messages.length = 0;
+  await passesItemLocator.hover();
+  await passesItemLocator.getByTitle('Open in VS Code').click();
+
+  expect(messages).toEqual([{
+    method: 'open',
+    params: {
+      location: expect.stringContaining('a.test.ts:5'),
+    },
+  }]);
+
+  await expect(
+      page.getByTestId('source-code').locator('.source-tab-file-name')
+  ).toHaveText('a.test.ts');
+  await expect(page.locator('.CodeMirror-code')).toContainText(`3      test('new-test', () => {});`);
+});