Add tests for select location entry and exit relay lists

2 files changed, 89 insertions, 2 deletions

diff --git a/desktop/packages/mullvad-vpn/test/e2e/mocked/select-location/helpers.ts b/desktop/packages/mullvad-vpn/test/e2e/mocked/select-location/helpers.ts
index 283a35b48e..308bf95ad8 100644
--- a/desktop/packages/mullvad-vpn/test/e2e/mocked/select-location/helpers.ts
+++ b/desktop/packages/mullvad-vpn/test/e2e/mocked/select-location/helpers.ts
@@ -36,6 +36,16 @@ export const createHelpers = (page: Page, routes: RoutesObjectModel, utils: Mock
     }
   };
 
+  const locateRelaysByHostnames = (relayList: IRelayList, hostnames: string[]): LocatedRelay[] => {
+    return relayList.countries.flatMap((country) =>
+      country.cities.flatMap((city) =>
+        city.relays
+          .filter((relay) => hostnames.includes(relay.hostname))
+          .map((relay) => ({ country, city, relay })),
+      ),
+    );
+  };
+
   const locateRelaysByProvider = (relayList: IRelayList, provider?: string): LocatedRelay[] =>
     relayList.countries.flatMap((country) =>
       country.cities.flatMap((city) =>
@@ -176,6 +186,7 @@ export const createHelpers = (page: Page, routes: RoutesObjectModel, utils: Mock
   return {
     areAllCheckboxesChecked,
     expandLocatedRelays,
+    locateRelaysByHostnames,
     locateRelaysByProvider,
     locateRelaysByOwner,
     locateRelaysByObfuscation,
diff --git a/desktop/packages/mullvad-vpn/test/e2e/mocked/select-location/select-location.spec.ts b/desktop/packages/mullvad-vpn/test/e2e/mocked/select-location/select-location.spec.ts
index 26c3e320ca..38aa4a3b2e 100644
--- a/desktop/packages/mullvad-vpn/test/e2e/mocked/select-location/select-location.spec.ts
+++ b/desktop/packages/mullvad-vpn/test/e2e/mocked/select-location/select-location.spec.ts
@@ -25,8 +25,12 @@ test.describe('Select location', () => {
     await helpers.updateMockRelays(relayList);
 
     await util.waitForRoute(RoutePath.main);
-    await page.getByLabel('Select location').click();
-    await util.waitForRoute(RoutePath.selectLocation);
+  });
+
+  test.beforeEach(async () => {
+    if ((await util.currentRoute()) === RoutePath.main) {
+      await routes.main.gotoSelectLocation();
+    }
   });
 
   test.afterAll(async () => {
@@ -88,6 +92,78 @@ test.describe('Select location', () => {
       const sweden = page.getByText('Sweden');
       await expect(sweden).toBeVisible();
     });
+
+    test('Should show only wireguard servers in entry list', async () => {
+      const entryButton = routes.selectLocation.getEntryButton();
+      await entryButton.click();
+
+      const wireguardRelays = relayList.countries[0].cities[0].relays.filter(
+        (relay) => relay.endpointType === 'wireguard',
+      );
+      const hostnames = wireguardRelays.map((relay) => relay.hostname);
+      const locatedRelays = helpers.locateRelaysByHostnames(relayList, hostnames);
+
+      await helpers.expandLocatedRelays(locatedRelays);
+
+      const buttons = routes.selectLocation.getRelaysMatching(hostnames);
+      await expect(buttons).toHaveCount(wireguardRelays.length);
+    });
+
+    test('Should show only wireguard servers in exit list', async () => {
+      const exitButton = routes.selectLocation.getExitButton();
+      await exitButton.click();
+
+      const wireguardRelays = relayList.countries[0].cities[0].relays.filter(
+        (relay) => relay.endpointType === 'wireguard',
+      );
+      const hostnames = wireguardRelays.map((relay) => relay.hostname);
+      const locatedRelays = helpers.locateRelaysByHostnames(relayList, hostnames);
+
+      await helpers.expandLocatedRelays(locatedRelays);
+
+      const buttons = routes.selectLocation.getRelaysMatching(hostnames);
+      await expect(buttons).toHaveCount(wireguardRelays.length);
+    });
+
+    test('Should disable entry server in exit list', async () => {
+      const settings = await helpers.updateMockSettings({
+        multihop: true,
+        daita: true,
+        directOnly: true,
+      });
+
+      const entryButton = routes.selectLocation.getEntryButton();
+      await entryButton.click();
+
+      // Get first wireguard relay
+      const [entryRelay, exitRelay] = relayList.countries[0].cities[0].relays.filter(
+        (relay) => relay.endpointType === 'wireguard',
+      );
+
+      if (!entryRelay) {
+        throw new Error('No wireguard relay found in mocked data');
+      }
+
+      const locatedEntryRelay = helpers.locateRelaysByHostnames(relayList, [entryRelay.hostname]);
+
+      await helpers.expandLocatedRelays(locatedEntryRelay);
+
+      await routes.selectLocation.getRelaysMatching([entryRelay.hostname]).first().click();
+
+      await helpers.updateEntryLocation(locatedEntryRelay[0], settings);
+
+      await helpers.expandLocatedRelays(locatedEntryRelay);
+      const entryRelayButton = routes.selectLocation.getRelaysMatching([entryRelay.hostname]);
+      await expect(entryRelayButton).toBeDisabled();
+
+      const locatedExitRelay = helpers.locateRelaysByHostnames(relayList, [exitRelay.hostname]);
+      await helpers.expandLocatedRelays(locatedExitRelay);
+
+      // Clicking exit relay should navigate to main route
+      const exitRelayButton = routes.selectLocation.getRelaysMatching([exitRelay.hostname]);
+      await exitRelayButton.click();
+      await util.waitForRoute(RoutePath.main);
+    });
   });
 
   test.describe('Filter', () => {