Fix some tests

nikhilwoodruff · nikhilwoodruff · commit a0de606cc076 · 2025-05-26T18:59:30.000+01:00
diff --git a/.github/workflows/any_changes.yaml b/.github/workflows/any_changes.yaml
@@ -26,8 +26,6 @@ jobs:
         
         - name: Test documentation builds
           run: make documentation
-          env:
-            HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
 
         - name: Check documentation build
           run: |
diff --git a/.github/workflows/code_changes.yaml b/.github/workflows/code_changes.yaml
@@ -43,6 +43,4 @@ jobs:
             run: uv pip install .[dev] --system
           
           - name: Run tests
-            run: make test
-            env:
-              HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
+            run: make test
diff --git a/.github/workflows/publish_documentation.yaml b/.github/workflows/publish_documentation.yaml
@@ -29,8 +29,6 @@ jobs:
         
       - name: Build documentation
         run: make documentation
-        env:
-          HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
       
       - name: Deploy documentation
         uses: JamesIves/github-pages-deploy-action@releases/v3
diff --git a/.github/workflows/publish_package.yaml b/.github/workflows/publish_package.yaml
@@ -36,8 +36,6 @@ jobs:
         
       - name: Test documentation builds
         run: make documentation
-        env:
-          HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
       
       - name: Deploy documentation
         uses: JamesIves/github-pages-deploy-action@releases/v3
diff --git a/.gitignore b/.gitignore
@@ -162,3 +162,5 @@ cython_debug/
 *.ipynb
 
 !docs/**/*.ipynb
+
+**/*.h5
diff --git a/policyengine/outputs/macro/comparison/calculate_economy_comparison.py b/policyengine/outputs/macro/comparison/calculate_economy_comparison.py
@@ -711,7 +711,6 @@ def uk_constituency_breakdown(
     reform_hnet = reform.household_net_income
 
     constituency_weights_path = download(
-        huggingface_repo="policyengine-uk-data",
         gcs_bucket="policyengine-uk-data-private",
         filepath="parliamentary_constituency_weights.h5",
     )
@@ -721,7 +720,6 @@ def uk_constituency_breakdown(
         ]  # {2025: array(650, 100180) where cell i, j is the weight of household record i in constituency j}
 
     constituency_names_path = download(
-        huggingface_repo="policyengine-uk-data",
         gcs_bucket="policyengine-uk-data-private",
         filepath="constituencies_2024.csv",
     )
diff --git a/policyengine/simulation.py b/policyengine/simulation.py
@@ -67,6 +67,10 @@ class SimulationOptions(BaseModel):
         None,
         description="The version of the country model used in the simulation. If not provided, the current package version will be used. If provided, this package will throw an error if the package version does not match. Use this as an extra safety check.",
     )
+    data_version: str | None = Field(
+        None,
+        description="The version of the data used in the simulation. If not provided, the current data version will be used. If provided, this package will throw an error if the data version does not match. Use this as an extra safety check.",
+    )
 
 
 class Simulation:
@@ -133,12 +137,7 @@ def _set_data(self):
                 bucket, filename = self.options.data.split("://")[-1].split(
                     "/"
                 )
-
-                if "@" in filename:
-                    filename, version = filename.split("@")
-                    self.data_version = version
-                else:
-                    version = None
+                version = self.options.data_version
 
                 file_path = download(
                     filepath=filename,
@@ -267,7 +266,6 @@ def _apply_region_to_simulation(
             elif "constituency/" in region:
                 constituency = region.split("/")[1]
                 constituency_names_file_path = download(
-                    huggingface_repo="policyengine-uk-data",
                     gcs_bucket="policyengine-uk-data-private",
                     filepath="constituencies_2024.csv",
                 )
@@ -288,7 +286,6 @@ def _apply_region_to_simulation(
                         f"Constituency {constituency} not found. See {constituency_names_file_path} for the list of available constituencies."
                     )
                 weights_file_path = download(
-                    huggingface_repo="policyengine-uk-data",
                     gcs_bucket="policyengine-uk-data-private",
                     filepath="parliamentary_constituency_weights.h5",
                 )
@@ -304,7 +301,6 @@ def _apply_region_to_simulation(
             elif "local_authority/" in region:
                 la = region.split("/")[1]
                 la_names_file_path = download(
-                    huggingface_repo="policyengine-uk-data",
                     gcs_bucket="policyengine-uk-data-private",
                     filepath="local_authorities_2021.csv",
                 )
@@ -319,7 +315,6 @@ def _apply_region_to_simulation(
                         f"Local authority {la} not found. See {la_names_file_path} for the list of available local authorities."
                     )
                 weights_file_path = download(
-                    huggingface_repo="policyengine-uk-data",
                     gcs_bucket="policyengine-uk-data-private",
                     filepath="local_authority_weights.h5",
                 )
diff --git a/policyengine/utils/data/caching_google_storage_client.py b/policyengine/utils/data/caching_google_storage_client.py
@@ -21,20 +21,6 @@ def __init__(self):
     def _data_key(self, bucket: str, key: str, version: str | None) -> str:
         return f"{bucket}.{key}.{version}.data"
 
-    def _get_latest_version(self, bucket: str, key: str) -> str | None:
-        """
-        Get the latest version of a blob in the specified bucket and key.
-        If no version is specified, return None.
-        """
-        blob = self.client.client.get_bucket(bucket).get_blob(key)
-        if blob.metadata is None:
-            logging.warning(
-                "No metadata found for blob, so it has no version attached."
-            )
-            return None
-        else:
-            return blob.metadata.get("version")
-
     # To absolutely 100% avoid any possible issue with file corruption or thread contention
     # always replace the current target file with whatever we have cached as an atomic write.
     def download(
@@ -45,7 +31,7 @@ def download(
         """
         if version is None:
             # If no version is specified, get the latest version from the cache
-            version = self._get_latest_version(bucket, key)
+            version = self.client._get_latest_version(bucket, key)
             logging.warning(
                 f"No version specified for {bucket}, {key}. Using latest version: {version}"
             )
diff --git a/policyengine/utils/data/simplified_google_storage_client.py b/policyengine/utils/data/simplified_google_storage_client.py
@@ -67,3 +67,17 @@ def download(
         # downloading the content. As a result this should now be the crc of the downloaded
         # content (i.e. there is not a race condition where it's getting the CRC from the cloud)
         return (result, blob.crc32c)
+
+    def _get_latest_version(self, bucket: str, key: str) -> str | None:
+        """
+        Get the latest version of a blob in the specified bucket and key.
+        If no version is specified, return None.
+        """
+        blob = self.client.get_bucket(bucket).get_blob(key)
+        if blob.metadata is None:
+            logging.warning(
+                "No metadata found for blob, so it has no version attached."
+            )
+            return None
+        else:
+            return blob.metadata.get("version")
diff --git a/policyengine/utils/data_download.py b/policyengine/utils/data_download.py
@@ -1,7 +1,6 @@
 from pathlib import Path
 import logging
 import os
-from policyengine.utils.huggingface import download_from_hf
 from policyengine.utils.google_cloud_bucket import download_file_from_gcs
 from pydantic import BaseModel
 import json
diff --git a/policyengine/utils/huggingface.py b/policyengine/utils/huggingface.py
diff --git a/tests/country/test_uk.py b/tests/country/test_uk.py
@@ -33,9 +33,7 @@ def test_uk_macro_bad_package_versions_fail():
             reform={
                 "gov.hmrc.income_tax.allowances.personal_allowance.amount": 15_000,
             },
-            package_versions={
-                "policyengine-uk": "-1.0.0",
-            },
+            model_version="a",
         )
         raise ValueError(
             "Simulation should have failed with a bad package version."
@@ -44,7 +42,7 @@ def test_uk_macro_bad_package_versions_fail():
         pass
 
 
-def test_uk_macro_bad_data_versions_fail():
+def test_uk_macro_bad_data_version_fails():
     from policyengine import Simulation
 
     try:
@@ -54,9 +52,7 @@ def test_uk_macro_bad_data_versions_fail():
             reform={
                 "gov.hmrc.income_tax.allowances.personal_allowance.amount": 15_000,
             },
-            data_versions={
-                "enhanced_frs_2022_23.h5": "-1.0.0",
-            },
+            data_version="a",
         )
         raise ValueError(
             "Simulation should have failed with a bad data version."
diff --git a/tests/utils/data/conftest.py b/tests/utils/data/conftest.py
@@ -12,6 +12,10 @@ def given_stored_data(self, data: str, crc: str):
             data.encode(),
             crc,
         )
+        print("Setting latest version")
+        self.mock_simple_storage_client._get_latest_version.return_value = (
+            "1.2.3"
+        )
 
     def given_crc_changes_on_download(
         self, data: str, initial_crc: str, download_crc: str
@@ -21,6 +25,9 @@ def given_crc_changes_on_download(
             data.encode(),
             download_crc,
         )
+        self.mock_simple_storage_client._get_latest_version.return_value = (
+            "1.2.3"
+        )
 
 
 @pytest.fixture()
diff --git a/tests/utils/data/test_google_cloud_bucket.py b/tests/utils/data/test_google_cloud_bucket.py
@@ -19,10 +19,13 @@ def setUp(self):
     def test_download_uses_storage_client(self, client_class):
         client_instance = client_class.return_value
         download_file_from_gcs(
-            "TEST_BUCKET", "TEST/FILE/NAME.TXT", "TARGET/PATH"
+            "TEST_BUCKET", "TEST/FILE/NAME.TXT", "TARGET/PATH", version=None
         )
         client_instance.download.assert_called_with(
-            "TEST_BUCKET", "TEST/FILE/NAME.TXT", Path("TARGET/PATH")
+            "TEST_BUCKET",
+            "TEST/FILE/NAME.TXT",
+            Path("TARGET/PATH"),
+            version=None,
         )
 
     @patch(
@@ -31,9 +34,9 @@ def test_download_uses_storage_client(self, client_class):
     )
     def test_download_only_creates_client_once(self, client_class):
         download_file_from_gcs(
-            "TEST_BUCKET", "TEST/FILE/NAME.TXT", "TARGET/PATH"
+            "TEST_BUCKET", "TEST/FILE/NAME.TXT", "TARGET/PATH", version=None
         )
         download_file_from_gcs(
-            "TEST_BUCKET", "TEST/FILE/NAME.TXT", "ANOTHER/PATH"
+            "TEST_BUCKET", "TEST/FILE/NAME.TXT", "ANOTHER/PATH", version=None
         )
         client_class.assert_called_once()

-Original file line number
+Diff line change
 *.ipynb
 !docs/**/*.ipynb
++
 +**/*.h5
Original file line number	Diff line number	Diff line change
`@@ -711,7 +711,6 @@ def uk_constituency_breakdown(`
`711`	`711`	`reform_hnet = reform.household_net_income`
`712`	`712`
`713`	`713`	`constituency_weights_path = download(`
`714`		`- huggingface_repo="policyengine-uk-data",`
`715`	`714`	`gcs_bucket="policyengine-uk-data-private",`
`716`	`715`	`filepath="parliamentary_constituency_weights.h5",`
`717`	`716`	`)`
`@@ -721,7 +720,6 @@ def uk_constituency_breakdown(`
`721`	`720`	`] # {2025: array(650, 100180) where cell i, j is the weight of household record i in constituency j}`
`722`	`721`
`723`	`722`	`constituency_names_path = download(`
`724`		`- huggingface_repo="policyengine-uk-data",`
`725`	`723`	`gcs_bucket="policyengine-uk-data-private",`
`726`	`724`	`filepath="constituencies_2024.csv",`
`727`	`725`	`)`