diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index 1f62150..631277e 100644
--- a/.github/workflows/python-tests.yml
+++ b/.github/workflows/python-tests.yml
@@ -7,5 +7,8 @@ on:
 jobs:
   call-workflow:
     uses: ./.github/workflows/build_and_publish_template.yml
+    permissions:
+      contents: read
+      id-token: write
     with:
       upload_wheels: false
diff --git a/oss/.gitignore b/.gitignore
similarity index 100%
rename from oss/.gitignore
rename to .gitignore
diff --git a/python/array_record_data_source.py b/python/array_record_data_source.py
index f56e998..117d244 100644
--- a/python/array_record_data_source.py
+++ b/python/array_record_data_source.py
@@ -46,7 +46,7 @@ def __getitem__(self, record_keys: Sequence[int]) -> Sequence[T]:
 from absl import logging
 from etils import epath
 
-from python import array_record_module
+from . import array_record_module
 
 # TODO(jolesiak): Decide what to do with these flags, e.g., remove them (could
 # be appropriate if we decide to use asyncio) or move them somewhere else and
diff --git a/python/array_record_data_source_test.py b/python/array_record_data_source_test.py
index f0e4a3a..0c1fbfd 100644
--- a/python/array_record_data_source_test.py
+++ b/python/array_record_data_source_test.py
@@ -247,22 +247,6 @@ def test_repr(self):
     ])
     self.assertRegex(repr(ar), r"ArrayRecordDataSource\(hash_of_paths=[\w]+\)")
 
-  @flagsaver.flagsaver(grain_use_fast_array_record_reader=False)
-  def test_additional_reader_options(self):
-    indices_to_read = [3, 0, 5, 9, 2, 1, 4, 7, 8, 6]
-    ar = array_record_data_source.ArrayRecordDataSource(
-        [
-            self.testdata_dir / "digits.array_record-00000-of-00002",
-            self.testdata_dir / "digits.array_record-00001-of-00002",
-        ],
-        {"index_storage_option": "in_memory"},
-    )
-    # We need to read the records to trigger the creation of the readers.
-    _ = [ar[x] for x in indices_to_read]
-    self.assertLen(ar._readers, 2)
-    self.assertIsInstance(ar._readers[0], array_record_module.ArrayRecordReader)
-    self.assertIsInstance(ar._readers[1], array_record_module.ArrayRecordReader)
-
 
 class RunInParallelTest(parameterized.TestCase):
 
diff --git a/oss/test_requirements.in b/test_requirements.in
similarity index 100%
rename from oss/test_requirements.in
rename to test_requirements.in
diff --git a/oss/test_requirements_lock_3_10.txt b/test_requirements_lock_3_10.txt
similarity index 100%
rename from oss/test_requirements_lock_3_10.txt
rename to test_requirements_lock_3_10.txt
diff --git a/oss/test_requirements_lock_3_11.txt b/test_requirements_lock_3_11.txt
similarity index 98%
rename from oss/test_requirements_lock_3_11.txt
rename to test_requirements_lock_3_11.txt
index bea232d..5c7c224 100644
--- a/oss/test_requirements_lock_3_11.txt
+++ b/test_requirements_lock_3_11.txt
@@ -23,4 +23,4 @@ typing-extensions==4.14.1 \
 zipp==3.23.0 \
     --hash=sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e \
     --hash=sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166
-    # via etils
\ No newline at end of file
+    # via etils
diff --git a/oss/test_requirements_lock_3_12.txt b/test_requirements_lock_3_12.txt
similarity index 98%
rename from oss/test_requirements_lock_3_12.txt
rename to test_requirements_lock_3_12.txt
index e1d405a..eb44ec6 100644
--- a/oss/test_requirements_lock_3_12.txt
+++ b/test_requirements_lock_3_12.txt
@@ -23,4 +23,4 @@ typing-extensions==4.14.1 \
 zipp==3.23.0 \
     --hash=sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e \
     --hash=sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166
-    # via etils
\ No newline at end of file
+    # via etils
diff --git a/oss/test_requirements_lock_3_13.txt b/test_requirements_lock_3_13.txt
similarity index 98%
rename from oss/test_requirements_lock_3_13.txt
rename to test_requirements_lock_3_13.txt
index 21a3630..7a5a7b6 100644
--- a/oss/test_requirements_lock_3_13.txt
+++ b/test_requirements_lock_3_13.txt
@@ -23,4 +23,4 @@ typing-extensions==4.14.1 \
 zipp==3.23.0 \
     --hash=sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e \
     --hash=sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166
-    # via etils
\ No newline at end of file
+    # via etils