check url, add resstock test

tests/models/test_models.py

@@ -4,17 +4,19 @@
 import os
 import shutil
 
-def test_model_downloads():
 
+def test_model_downloads_comstock():
+
+    folder_name = os.path.join(os.getcwd(), 'comsstock')
     objA = Inputs
-    objA.__init__(Inputs, '2023','comstock_amy2018_release_2','18')
+    objA.__init__(Inputs, '2023','comstock_amy2018_release_2','18', folder_name)
 
     id = ['bldg0000001','bldg0000002']
-    folder_name = 'saved'
+
     shutil.rmtree(folder_name, ignore_errors=True)
 
     objB = Models
-    objB.download_models(id, objA, folder_name)
+    objB.download_models(id, objA)
 
 
     save_folder_1 = os.path.join(folder_name, id[0])
@@ -24,3 +26,21 @@ def test_model_downloads():
     assert os.path.exists(save_folder_1)
     assert os.path.exists(save_folder_2)
 
+def test_model_downloads_resstock():
+
+    folder_name = os.path.join(os.getcwd(), 'resstock')
+    objA = Inputs
+    objA.__init__(Inputs, '2021','resstock_tmy3_release_1','00', folder_name)
+
+    id = ['bldg0000003','bldg0000004']
+    shutil.rmtree(folder_name, ignore_errors=True)
+
+    objB = Models
+    objB.download_models(id, objA)
+
+    save_folder_1 = os.path.join(folder_name, id[0])
+    save_folder_2 = os.path.join(folder_name, id[1])
+
+    assert os.path.exists(folder_name)
+    assert os.path.exists(save_folder_1)
+    assert os.path.exists(save_folder_2)

workflow/lib/inputs.py

@@ -4,7 +4,8 @@
 class Inputs:
     """Class to assign inputs for running buildstock models in Alfalfa."""
 
-    def __init__(self, year, dataset_name, upgrade):
+    def __init__(self, year, dataset_name, upgrade, folder):
         self.year = year
         self.dataset_name = dataset_name
         self.upgrade = upgrade
+        self.folder = folder

workflow/lib/models.py

@@ -4,24 +4,29 @@
 import gzip
 import os
 from workflow.lib.inputs import Inputs
-from urllib.parse import urljoin, urlsplit, urlunsplit
+from urllib.parse import urljoin, urlsplit, urlunsplit, urlparse
 from pathlib import Path
 from io import BytesIO
 
-
 class Models:
 
     """Class to download models from defined inputs"""
 
-    def download_models(id, obj, folder_path):
+    def download_models(id, obj):
 
+        folder_path = obj.folder
         os.makedirs(folder_path, exist_ok=True)  # Create folder if it doesn't exist
 
         for i in id:
             # TO DO: check if this url is different for other years/resstock
-            osm_url = f'https://oedi-data-lake.s3.amazonaws.com/nrel-pds-building-stock/end-use-load-profiles-for-us-building-stock/{obj.year}/{obj.dataset_name}/building_energy_models/upgrade={obj.upgrade}/{i}-up{obj.upgrade}.osm.gz'
-            response = requests.get(osm_url, stream=True)
+            url1 = f'https://oedi-data-lake.s3.amazonaws.com/nrel-pds-building-stock/end-use-load-profiles-for-us-building-stock/{obj.year}/{obj.dataset_name}/building_energy_models/upgrade={obj.upgrade}/{i}-up{obj.upgrade}.osm.gz'
+            url2 = f'https://oedi-data-lake.s3.amazonaws.com/nrel-pds-building-stock/end-use-load-profiles-for-us-building-stock/{obj.year}/{obj.dataset_name}/building_energy_models/{i}-up{obj.upgrade}.osm.gz'
+
+            valid_url = next((url for url in (url1, url2) if requests.head(url).status_code == 200), None)
+            if not valid_url:
+                raise ValueError("Invalid URL")
 
+            response = requests.get(valid_url, stream= True)
             save_folder = os.path.join(folder_path, i)
             os.makedirs(save_folder, exist_ok=True)  # Create folder using filename if it doesn't exist
             osm_path = os.path.join(save_folder, i + '.osm')