From d43d2822c0a8c7e021bf2185db771e6407df5314 Mon Sep 17 00:00:00 2001
From: Xee authors <no-reply@google.com>
Date: Mon, 6 Nov 2023 08:24:07 -0800
Subject: [PATCH] No public description

PiperOrigin-RevId: 579850276
---
 .github/workflows/ci-build.yml | 17 +++++++++++++----
 setup.py                       |  1 +
 xee/ext_integration_test.py    | 24 ++++++++++++++++++++++--
 3 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/ci-build.yml b/.github/workflows/ci-build.yml
index 2e761be..6d7e746 100644
--- a/.github/workflows/ci-build.yml
+++ b/.github/workflows/ci-build.yml
@@ -31,6 +31,8 @@ jobs:
       fail-fast: false
       matrix:
         python-version: ["3.9", "3.10", "3.11"]
+    permissions:
+      id-token: write  # This is required for requesting the JWT.
     steps:
     - name: Cancel previous
       uses: styfle/cancel-workflow-action@0.7.0
@@ -55,10 +57,17 @@ jobs:
     - name: Install Xee
       run: |
         pip install -e .[tests]
+    # Fix an issue with "Unrecognized chunk manager. Must be one of []."
+    - name: Fix xarray installation
+      run: |
+        pip install --force-reinstall "xarray==0.21.1"
+    - uses: 'actions/checkout@v4'
+    - id: 'auth'
+      name: 'Authenticate to Google Cloud'
+      uses: 'google-github-actions/auth@v1'
+      with:
+        service_account: ${{ secrets.SERVICE_ACCOUNT }}
+        workload_identity_provider: ${{ secrets.WORKLOAD_IDENTITY_PROVIDER }}
     - name: Run unit tests
       run: |
         pytest xee
-    - name: Run scripts tests
-      # The scripts may define some of the same flags, so we run pytest in separate processes.
-      run: |
-        for test in examples/*_test.py; do pytest $test; done
diff --git a/setup.py b/setup.py
index 9d8a433..de9b4b0 100644
--- a/setup.py
+++ b/setup.py
@@ -26,6 +26,7 @@
 
 tests_requires = [
     "absl-py",
+    "dask",
     "pytest",
     "pyink",
 ]
diff --git a/xee/ext_integration_test.py b/xee/ext_integration_test.py
index 8e36d4d..e68f100 100644
--- a/xee/ext_integration_test.py
+++ b/xee/ext_integration_test.py
@@ -13,9 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 r"""Integration tests for the Google Earth Engine backend for Xarray."""
+import json
+import os
 import pathlib
 
 from absl.testing import absltest
+from google.auth import identity_pool
 import numpy as np
 import xarray as xr
 from xarray.core import indexing
@@ -23,9 +26,26 @@
 
 import ee
 
+_CREDENTIALS_PATH_KEY = 'GOOGLE_APPLICATION_CREDENTIALS'
+_SCOPES = [
+    'https://www.googleapis.com/auth/cloud-platform',
+    'https://www.googleapis.com/auth/earthengine',
+]
+
+
+def _read_identity_pool_creds() -> identity_pool.Credentials:
+  credentials_path = os.environ[_CREDENTIALS_PATH_KEY]
+  with open(credentials_path) as file:
+    json_file = json.load(file)
+    credentials = identity_pool.Credentials.from_info(json_file)
+    return credentials.with_scopes(_SCOPES)
+
 
 def init_ee_for_tests():
-  ee.Initialize(opt_url='https://earthengine-highvolume.googleapis.com')
+  ee.Initialize(
+      credentials=_read_identity_pool_creds(),
+      opt_url='https://earthengine-highvolume.googleapis.com',
+  )
 
 
 class EEBackendArrayTest(absltest.TestCase):
@@ -337,7 +357,7 @@ def test_parses_ee_url(self):
         scale=25.0,  # in degrees
         n_images=3,
     )
-    self.assertEqual(dict(ds.dims), {'time': 3, 'lon': 15, 'lat': 7})
+    self.assertEqual(dict(ds.dims), {'time': 3, 'lon': 15, 'lat': 8})
     ds = self.entry.open_dataset(
         'ee:LANDSAT/LC08/C01/T1',
         drop_variables=tuple(f'B{i}' for i in range(3, 12)),