GAIAMNGT-1700 New tests

Author: Jorge Fernandez Hernandez

astroquery/gaia/core.py

@@ -218,7 +218,8 @@ def load_data(self, ids, *, data_release=None, data_structure='INDIVIDUAL', retr
         format : str, optional, default 'votable'
             loading format. Other available formats are 'csv', 'ecsv','votable_plain', 'json' and 'fits'
         dump_to_file: boolean, optional, default False.
-            If it is true, a compressed directory named "datalink_output.zip" with all the DataLink files is made
+            If it is true, a compressed directory named "datalink_output.zip" with all the DataLink files is made in the
+             current working directory
         overwrite_output_file : boolean, optional, default False
             To overwrite the output_file if it already exists.
         verbose : bool, optional, default 'False'
@@ -242,22 +243,23 @@ def load_data(self, ids, *, data_release=None, data_structure='INDIVIDUAL', retr
             output_file = 'datalink_output.zip'
             output_file_specified = True
             output_file = os.path.abspath(output_file)
-            print(f"DataLink products are stored inside the {output_file} file")
+            log.info(f"DataLink products will be stored in the {output_file} file")
 
             if not overwrite_output_file and os.path.exists(output_file):
-                print(f"{output_file} file already exists and will be overwritten")
+                log.warn(f"{output_file} file already exists and will be overwritten")
 
         path = os.path.dirname(output_file)
 
         log.debug(f"Directory where the data will be saved: {path}")
 
         if path != '':
-            try:
-                os.mkdir(path)
-            except FileExistsError:
-                log.error("Path %s already exist" % path)
-            except OSError:
-                log.error("Creation of the directory %s failed" % path)
+            if not os.path.isdir(path):
+                try:
+                    os.mkdir(path)
+                except FileExistsError:
+                    log.warn("Path %s already exist" % path)
+                except OSError:
+                    log.error("Creation of the directory %s failed" % path)
 
         if avoid_datatype_check is False:
             # we need to check params
@@ -312,8 +314,7 @@ def load_data(self, ids, *, data_release=None, data_structure='INDIVIDUAL', retr
                 shutil.rmtree(path)
             else:
                 for file in files.keys():
-                    os.remove(os.path.join(os.getcwd(), path, file)
-                              )
+                    os.remove(os.path.join(os.getcwd(), path, file))
 
         if verbose:
             if output_file_specified:
@@ -332,9 +333,8 @@ def __get_data_files(output_file, path):
         extracted_files = []
 
         with zipfile.ZipFile(output_file, "r") as zip_ref:
-            for name in zip_ref.namelist():
-                local_file_path = zip_ref.extract(name, os.path.dirname(output_file))
-                extracted_files.append(local_file_path)
+            extracted_files.extend(zip_ref.namelist())
+            zip_ref.extractall(os.path.dirname(output_file))
 
         # r=root, d=directories, f = files
         for r, d, f in os.walk(path):
@@ -344,7 +344,7 @@ def __get_data_files(output_file, path):
 
         for key, value in files.items():
 
-            if '.fits' in key:
+            if key.endswith('.fits'):
                 tables = []
                 with fits.open(value) as hduList:
                     num_hdus = len(hduList)
@@ -354,19 +354,19 @@ def __get_data_files(output_file, path):
                         tables.append(table)
                     files[key] = tables
 
-            elif '.xml' in key:
+            elif key.endswith('.xml'):
                 tables = []
                 for table in votable.parse(value).iter_tables():
                     tables.append(table)
                 files[key] = tables
 
-            elif '.csv' in key:
+            elif key.endswith('.csv'):
                 tables = []
                 table = Table.read(value, format='ascii.csv', fast_reader=False)
                 tables.append(table)
                 files[key] = tables
 
-            elif '.json' in key:
+            elif key.endswith('.json'):
                 tables = []
                 with open(value) as f:
                     data = json.load(f)

astroquery/gaia/tests/data/gaia_dr3_source_id_5937083312263887616_dl_products_csv.zip

@@ -10,7 +10,8 @@ def get_package_data():
     paths = [os.path.join('data', '*.vot'),
              os.path.join('data', '*.vot.gz'),
              os.path.join('data', '*.json'),
-             os.path.join('data', '*.ecsv')
+             os.path.join('data', '*.ecsv'),
+             os.path.join('data', '*.zip')
              ]  # etc, add other extensions
     # you can also enlist files individually by names
     # finally construct and return a dict for the sub module

astroquery/gaia/tests/data/gaia_dr3_source_id_5937083312263887616_dl_products_ecsv.zip

@@ -15,6 +15,7 @@
 
 """
 import os
+import zipfile
 from pathlib import Path
 from unittest.mock import patch
 
@@ -45,6 +46,22 @@
 JOB_DATA_QUERIER_ECSV_FILE_NAME = get_pkg_data_filename(os.path.join("data", '1712337806100O-result.ecsv'),
                                                         package=package)
 
+DL_PRODUCTS_VOT = get_pkg_data_filename(
+    os.path.join("data", 'gaia_dr3_source_id_5937083312263887616_dl_products_vot.zip'),
+    package=package)
+
+DL_PRODUCTS_ECSV = get_pkg_data_filename(
+    os.path.join("data", 'gaia_dr3_source_id_5937083312263887616_dl_products_ecsv.zip'),
+    package=package)
+
+DL_PRODUCTS_CSV = get_pkg_data_filename(
+    os.path.join("data", 'gaia_dr3_source_id_5937083312263887616_dl_products_csv.zip'),
+    package=package)
+
+DL_PRODUCTS_FITS = get_pkg_data_filename(
+    os.path.join("data", 'gaia_dr3_source_id_5937083312263887616_dl_products_fits.zip'),
+    package=package)
+
 JOB_DATA = Path(JOB_DATA_FILE_NAME).read_text()
 JOB_DATA_NEW = Path(JOB_DATA_FILE_NAME_NEW).read_text()
 
@@ -152,6 +169,23 @@ def mock_querier():
     return GaiaClass(tap_plus_conn_handler=conn_handler, datalink_handler=tapplus, show_server_messages=False)
 
 
+@pytest.fixture(scope="module")
+def mock_datalink_querier():
+    conn_handler = DummyConnHandler()
+    tapplus = TapPlus(url="http://test:1111/tap", connhandler=conn_handler)
+
+    launch_response = DummyResponse(200)
+    launch_response.set_data(method="POST", body=DL_PRODUCTS_VOT)
+    # The query contains decimals: default response is more robust.
+    conn_handler.set_default_response(launch_response)
+    conn_handler.set_response(
+        '?DATA_STRUCTURE=INDIVIDUAL&FORMAT=votable&ID=5937083312263887616&RELEASE=Gaia+DR3&RETRIEVAL_TYPE=ALL'
+        '&USE_ZIP_ALWAYS=true&VALID_DATA=false',
+        launch_response)
+
+    return GaiaClass(tap_plus_conn_handler=conn_handler, datalink_handler=tapplus, show_server_messages=False)
+
+
 @pytest.fixture(scope="module")
 def mock_querier_ecsv():
     conn_handler = DummyConnHandler()
@@ -696,29 +730,100 @@ def test_cone_search_and_changing_MAIN_GAIA_TABLE(mock_querier_async):
         assert "name_from_class" in job.parameters["query"]
 
 
-def test_load_data(monkeypatch, tmp_path):
+def test_load_data(mock_datalink_querier):
+    mock_datalink_querier.load_data(ids=[5937083312263887616], data_release='Gaia DR3', data_structure='INDIVIDUAL',
+                                    retrieval_type="ALL",
+                                    linking_parameter='SOURCE_ID', valid_data=False, band=None,
+                                    avoid_datatype_check=False,
+                                    format="votable", dump_to_file=True, overwrite_output_file=True, verbose=False)
+
+    assert os.path.exists('datalink_output.zip')
+
+    extracted_files = []
+
+    with zipfile.ZipFile('datalink_output.zip', "r") as zip_ref:
+        extracted_files.extend(zip_ref.namelist())
+
+    assert len(extracted_files) == 3
+
+    os.remove(os.path.join(os.getcwd(), 'datalink_output.zip'))
+
+    assert not os.path.exists('datalink_output.zip')
+
+
+@pytest.mark.skip(reason="Thes fits files generate an error relatate to the unit 'log(cm.s**-2)")
+def test_load_data_fits(monkeypatch, tmp_path, tmp_path_factory):
+    path = Path(os.getcwd() + '/' + 'datalink_output.zip')
+
+    with open(DL_PRODUCTS_FITS, 'rb') as file:
+        zip_bytes = file.read()
+
+    path.write_bytes(zip_bytes)
+
     def load_data_monkeypatched(self, params_dict, output_file, verbose):
         assert params_dict == {
             "VALID_DATA": "true",
             "ID": "1,2,3,4",
-            "FORMAT": "votable",
+            "FORMAT": "fits",
             "RETRIEVAL_TYPE": "epoch_photometry",
             "DATA_STRUCTURE": "INDIVIDUAL",
             "USE_ZIP_ALWAYS": "true"}
-        assert output_file == str(tmp_path / "datalink_output.zip")
+        assert output_file == os.getcwd() + '/' + 'datalink_output.zip'
         assert verbose is True
 
     monkeypatch.setattr(TapPlus, "load_data", load_data_monkeypatched)
 
     GAIA_QUERIER.load_data(
+        valid_data=True,
         ids="1,2,3,4",
+        format='fits',
         retrieval_type="epoch_photometry",
+        verbose=True,
+        dump_to_file=True)
+
+    path.unlink()
+
+
+def test_load_data_csv(monkeypatch, tmp_path, tmp_path_factory):
+    path = Path(os.getcwd() + '/' + 'datalink_output.zip')
+
+    with open(DL_PRODUCTS_CSV, 'rb') as file:
+        zip_bytes = file.read()
+
+    path.write_bytes(zip_bytes)
+
+    def load_data_monkeypatched(self, params_dict, output_file, verbose):
+        assert params_dict == {
+            "VALID_DATA": "true",
+            "ID": "1,2,3,4",
+            "FORMAT": "csv",
+            "RETRIEVAL_TYPE": "epoch_photometry",
+            "DATA_STRUCTURE": "INDIVIDUAL",
+            "USE_ZIP_ALWAYS": "true"}
+        assert output_file == os.getcwd() + '/' + 'datalink_output.zip'
+        assert verbose is True
+
+    monkeypatch.setattr(TapPlus, "load_data", load_data_monkeypatched)
+
+    GAIA_QUERIER.load_data(
         valid_data=True,
+        ids="1,2,3,4",
+        format='csv',
+        retrieval_type="epoch_photometry",
         verbose=True,
         dump_to_file=True)
 
+    path.unlink()
+
+
+def test_load_data_ecsv(monkeypatch, tmp_path, tmp_path_factory):
+    path = Path(os.getcwd() + '/' + 'datalink_output.zip')
+
+    with open(DL_PRODUCTS_ECSV, 'rb') as file:
+        zip_bytes = file.read()
+
+    path.write_bytes(zip_bytes)
 
-def test_load_data_ecsv(monkeypatch, tmp_path):
     def load_data_monkeypatched(self, params_dict, output_file, verbose):
         assert params_dict == {
             "VALID_DATA": "true",
@@ -727,21 +832,30 @@ def load_data_monkeypatched(self, params_dict, output_file, verbose):
             "RETRIEVAL_TYPE": "epoch_photometry",
             "DATA_STRUCTURE": "INDIVIDUAL",
             "USE_ZIP_ALWAYS": "true"}
-        assert output_file == str(tmp_path / "datalink_output.zip")
+        assert output_file == os.getcwd() + '/' + 'datalink_output.zip'
         assert verbose is True
 
     monkeypatch.setattr(TapPlus, "load_data", load_data_monkeypatched)
 
     GAIA_QUERIER.load_data(
+        valid_data=True,
         ids="1,2,3,4",
+        format='ecsv',
         retrieval_type="epoch_photometry",
-        valid_data=True,
         verbose=True,
-        format='ecsv',
         dump_to_file=True)
 
+    path.unlink()
+
 
 def test_load_data_linking_parameter(monkeypatch, tmp_path):
+    path = Path(os.getcwd() + '/' + 'datalink_output.zip')
+
+    with open(DL_PRODUCTS_VOT, 'rb') as file:
+        zip_bytes = file.read()
+
+    path.write_bytes(zip_bytes)
+
     def load_data_monkeypatched(self, params_dict, output_file, verbose):
         assert params_dict == {
             "VALID_DATA": "true",
@@ -750,7 +864,7 @@ def load_data_monkeypatched(self, params_dict, output_file, verbose):
             "RETRIEVAL_TYPE": "epoch_photometry",
             "DATA_STRUCTURE": "INDIVIDUAL",
             "USE_ZIP_ALWAYS": "true"}
-        assert output_file == str(tmp_path / "datalink_output.zip")
+        assert output_file == os.getcwd() + '/' + 'datalink_output.zip'
         assert verbose is True
 
     monkeypatch.setattr(TapPlus, "load_data", load_data_monkeypatched)
@@ -763,9 +877,18 @@ def load_data_monkeypatched(self, params_dict, output_file, verbose):
         verbose=True,
         dump_to_file=True)
 
+    path.unlink()
+
 
 @pytest.mark.parametrize("linking_param", ['TRANSIT_ID', 'IMAGE_ID'])
 def test_load_data_linking_parameter_with_values(monkeypatch, tmp_path, linking_param):
+    path = Path(os.getcwd() + '/' + 'datalink_output.zip')
+
+    with open(DL_PRODUCTS_VOT, 'rb') as file:
+        zip_bytes = file.read()
+
+    path.write_bytes(zip_bytes)
+
     def load_data_monkeypatched(self, params_dict, output_file, verbose):
         assert params_dict == {
             "VALID_DATA": "true",
@@ -775,7 +898,7 @@ def load_data_monkeypatched(self, params_dict, output_file, verbose):
             "DATA_STRUCTURE": "INDIVIDUAL",
             "LINKING_PARAMETER": linking_param,
             "USE_ZIP_ALWAYS": "true", }
-        assert output_file == str(tmp_path / "datalink_output.zip")
+        assert output_file == os.getcwd() + '/' + 'datalink_output.zip'
         assert verbose is True
 
     monkeypatch.setattr(TapPlus, "load_data", load_data_monkeypatched)
@@ -788,6 +911,8 @@ def load_data_monkeypatched(self, params_dict, output_file, verbose):
         verbose=True,
         dump_to_file=True)
 
+    path.unlink()
+
 
 def test_get_datalinks(monkeypatch):
     def get_datalinks_monkeypatched(self, ids, linking_parameter, verbose):

astroquery/gaia/tests/data/gaia_dr3_source_id_5937083312263887616_dl_products_fits.zip

@@ -24,6 +24,7 @@ class DummyResponse:
     STATUS_MESSAGES = {200: "OK", 303: "OK", 500: "ERROR"}
 
     def __init__(self, status_code=None):
+        self.zip_bytes = None
         self.reason = ""
         self.set_status_code(status_code)
         self.index = 0
@@ -57,11 +58,29 @@ def read(self, size=None):
         if v is None:
             return None
         else:
+
+            if v.endswith('zip'):
+                if self.zip_bytes is None:
+                    with open(v, 'rb') as file:
+                        self.zip_bytes = file.read()
+
             if size is None or size < 0:
+
+                if v.endswith('zip'):
+                    return self.zip_bytes
+
                 # read all
                 return v.encode(encoding='utf_8', errors='strict')
             else:
-                bodyLength = len(v)
+                is_zip = False
+
+                if v.endswith('zip'):
+                    is_zip = True
+                    bodyLength = len(self.zip_bytes)
+                    v = self.zip_bytes
+                else:
+                    bodyLength = len(v)
+
                 if self.index < 0:
                     return ""
                 if size >= bodyLength:
@@ -73,7 +92,11 @@ def read(self, size=None):
                 self.index = endPos
                 if endPos >= (bodyLength - 1):
                     self.index = -1
-                return tmp.encode(encoding='utf_8', errors='strict')
+
+                if is_zip:
+                    return tmp
+                else:
+                    return tmp.encode(encoding='utf_8', errors='strict')
 
     def close(self):
         self.index = 0