feat: geolocation controller fix

Author: MarioAslau

packages/geolocation-controller/src/GeolocationController.ts

@@ -22,7 +22,7 @@ export const controllerName = 'GeolocationController';
  * State for the {@link GeolocationController}.
  */
 export type GeolocationControllerState = {
-  /** ISO 3166-1 alpha-2 country code, or "UNKNOWN" if not yet determined. */
+  /** ISO 3166-2 location code (e.g. "US", "US-NY", "CA-ON"), or "UNKNOWN" if not yet determined. */
   location: string;
   /** Current status of the geolocation fetch lifecycle. */
   status: GeolocationRequestStatus;
@@ -182,11 +182,11 @@ export class GeolocationController extends BaseController<
   }
 
   /**
-   * Returns the geolocation country code. Delegates to the
+   * Returns the geolocation code. Delegates to the
    * {@link GeolocationApiService} for network fetching and caching, then
    * updates controller state with the result.
    *
-   * @returns The ISO country code string.
+   * @returns The ISO 3166-2 location code string.
    */
   async getGeolocation(): Promise<string> {
     return this.#fetchAndUpdate();
@@ -195,7 +195,7 @@ export class GeolocationController extends BaseController<
   /**
    * Forces a fresh geolocation fetch, bypassing the service's cache.
    *
-   * @returns The ISO country code string.
+   * @returns The ISO 3166-2 location code string.
    */
   async refreshGeolocation(): Promise<string> {
     this.update((draft) => {
@@ -210,7 +210,7 @@ export class GeolocationController extends BaseController<
    *
    * @param options - Options forwarded to the service.
    * @param options.bypassCache - When true, the service skips its TTL cache.
-   * @returns The ISO country code string.
+   * @returns The ISO 3166-2 location code string.
    */
   async #fetchAndUpdate(options?: { bypassCache?: boolean }): Promise<string> {
     this.update((draft) => {

packages/geolocation-controller/src/geolocation-api-service/geolocation-api-service.test.ts

@@ -246,7 +246,34 @@ describe('GeolocationApiService', () => {
         expect(result).toBe('US');
       });
 
-      it('returns UNKNOWN_LOCATION for non-ISO-3166-1 alpha-2 responses', async () => {
+      it('accepts a 2-letter country code', async () => {
+        const mockFetch = createMockFetch('GB');
+        const { service } = getService({ options: { fetch: mockFetch } });
+
+        const result = await service.fetchGeolocation();
+
+        expect(result).toBe('GB');
+      });
+
+      it('accepts a country code with subdivision (ISO 3166-2)', async () => {
+        const mockFetch = createMockFetch('US-NY');
+        const { service } = getService({ options: { fetch: mockFetch } });
+
+        const result = await service.fetchGeolocation();
+
+        expect(result).toBe('US-NY');
+      });
+
+      it('accepts a subdivision code with alphanumeric part', async () => {
+        const mockFetch = createMockFetch('CA-ON');
+        const { service } = getService({ options: { fetch: mockFetch } });
+
+        const result = await service.fetchGeolocation();
+
+        expect(result).toBe('CA-ON');
+      });
+
+      it('returns UNKNOWN_LOCATION for non-ISO responses', async () => {
         const mockFetch = jest
           .fn()
           .mockImplementation(() =>
@@ -272,7 +299,20 @@ describe('GeolocationApiService', () => {
         expect(result).toBe(UNKNOWN_LOCATION);
       });
 
-      it('returns UNKNOWN_LOCATION for three-letter codes', async () => {
+      it('returns UNKNOWN_LOCATION for lowercase subdivision codes', async () => {
+        const mockFetch = jest
+          .fn()
+          .mockImplementation(() =>
+            Promise.resolve(createMockResponse('us-ny', 200)),
+          );
+        const { service } = getService({ options: { fetch: mockFetch } });
+
+        const result = await service.fetchGeolocation();
+
+        expect(result).toBe(UNKNOWN_LOCATION);
+      });
+
+      it('returns UNKNOWN_LOCATION for three-letter country codes', async () => {
         const mockFetch = jest
           .fn()
           .mockImplementation(() =>
@@ -284,6 +324,19 @@ describe('GeolocationApiService', () => {
 
         expect(result).toBe(UNKNOWN_LOCATION);
       });
+
+      it('returns UNKNOWN_LOCATION for subdivision codes with too many characters', async () => {
+        const mockFetch = jest
+          .fn()
+          .mockImplementation(() =>
+            Promise.resolve(createMockResponse('US-ABCD', 200)),
+          );
+        const { service } = getService({ options: { fetch: mockFetch } });
+
+        const result = await service.fetchGeolocation();
+
+        expect(result).toBe(UNKNOWN_LOCATION);
+      });
     });
 
     describe('bypassCache', () => {
@@ -480,19 +533,19 @@ function createMockResponse(body: string, status: number): Response {
 }
 
 /**
- * Creates a mock fetch function that resolves with the given country code.
+ * Creates a mock fetch function that resolves with the given location code.
  * Each call returns a fresh mock Response.
  *
- * @param countryCode - The country code to return.
+ * @param locationCode - The location code to return (e.g. `US`, `US-NY`).
  * @returns A jest mock function.
  */
 function createMockFetch(
-  countryCode: string,
+  locationCode: string,
 ): jest.Mock<Promise<Response>, [string]> {
   return jest
     .fn()
     .mockImplementation(() =>
-      Promise.resolve(createMockResponse(countryCode, 200)),
+      Promise.resolve(createMockResponse(locationCode, 200)),
     );
 }
 

packages/geolocation-controller/src/geolocation-api-service/geolocation-api-service.ts

@@ -85,11 +85,12 @@ export type FetchGeolocationOptions = {
 };
 
 /**
- * Low-level data service that fetches a country code from the geolocation API.
+ * Low-level data service that fetches a location code from the geolocation API.
  *
  * Responsibilities:
  * - HTTP request to the geolocation endpoint (wrapped in a service policy)
- * - ISO 3166-1 alpha-2 response validation
+ * - ISO 3166-2 response validation (country code with optional subdivision,
+ *   e.g. `US`, `US-NY`, `CA-ON`)
  * - TTL-based in-memory cache
  * - Promise deduplication (concurrent callers share a single in-flight request)
  *
@@ -202,16 +203,16 @@ export class GeolocationApiService {
   }
 
   /**
-   * Returns the geolocation country code. Serves from cache when the TTL has
-   * not expired, otherwise performs a network fetch. Concurrent callers are
+   * Returns the geolocation code. Serves from cache when the TTL has not
+   * expired, otherwise performs a network fetch. Concurrent callers are
    * deduplicated to a single in-flight request.
    *
    * @param options - Optional fetch options.
    * @param options.bypassCache - When true, invalidates the TTL cache. If a
    * request is already in-flight it will be reused (deduplication always
    * applies).
-   * @returns The ISO 3166-1 alpha-2 country code, or {@link UNKNOWN_LOCATION}
-   * when the API returns an empty or invalid body.
+   * @returns An ISO 3166-2 location code (e.g. `US`, `US-NY`, `CA-ON`), or
+   * {@link UNKNOWN_LOCATION} when the API returns an empty or invalid body.
    */
   async fetchGeolocation(options?: FetchGeolocationOptions): Promise<string> {
     if (options?.bypassCache) {
@@ -252,7 +253,7 @@ export class GeolocationApiService {
    * Performs the actual HTTP fetch, wrapped in the service policy for automatic
    * retry and circuit-breaking, and validates the response.
    *
-   * @returns The ISO country code string.
+   * @returns The ISO 3166-2 location code string.
    */
   async #performFetch(): Promise<string> {
     const response = await this.#policy.execute(async () => {
@@ -267,7 +268,9 @@ export class GeolocationApiService {
     });
 
     const raw = (await response.text()).trim();
-    const location = /^[A-Z]{2}$/u.test(raw) ? raw : UNKNOWN_LOCATION;
+    const location = /^[A-Z]{2}(-[A-Z0-9]{1,3})?$/u.test(raw)
+      ? raw
+      : UNKNOWN_LOCATION;
 
     if (location !== UNKNOWN_LOCATION) {
       this.#cachedLocation = location;