Input validation to handle 502 Errors

imap_data_access/io.py

@@ -6,7 +6,9 @@
 import contextlib
 import json
 import logging
+import re
 import urllib.request
+from datetime import datetime
 from pathlib import Path
 from typing import Optional, Union
 from urllib.error import HTTPError, URLError
@@ -99,6 +101,8 @@ def download(file_path: Union[Path, str]) -> Path:
     return destination
 
 
+# Too many branches error
+# ruff: noqa: PLR0912
 def query(
     *,
     instrument: Optional[str] = None,
@@ -159,6 +163,85 @@ def query(
             "At least one query parameter must be provided. "
             "Run 'query -h' for more information."
         )
+    # Check instrument name
+    if instrument is not None:
+        if instrument not in (
+            "codice",
+            "glows",
+            "hi",
+            "hit",
+            "idex",
+            "lo",
+            "mag",
+            "swapi",
+            "swe",
+            "ultra",
+        ):
+            raise ValueError(
+                "Not a valid instrument, please choose from "
+                "('codice', 'glows', 'hi', 'hit', 'idex', "
+                "'lo', 'mag', 'swapi', 'swe', 'ultra')"
+            )
+
+    # Check data-level
+    # do an if statement that checks that data_level was passed in,
+    # then check it against all options, l0, l1a, l1b, l2, l3 etc.
+    if data_level is not None:
+        if data_level not in (
+            "l0",
+            "l1",
+            "l1a",
+            "l1b",
+            "l1c",
+            "l1ca",
+            "l1cb",
+            "l1d",
+            "l2",
+            "l2pre",
+            "l3",
+            "l3a",
+            "l3b",
+            "l3c",
+            "l3d",
+        ):
+            raise ValueError(
+                "Not a valid data level, choose from "
+                "('l0','l1','l1a','l1b','l1c','l1ca','l1cb',"
+                "'l1d','l2','l2pre','l3','l3a','l3b','l3c','l3d')."
+            )
+
+    # Check start-date
+    if start_date is not None:
+        try:
+            datetime.strptime(start_date, "%Y%m%d")
+        except ValueError as err:
+            raise ValueError("Not a valid start date, use format 'YYYYMMDD'.") from err
+
+    # Check end-date
+    if end_date is not None:
+        try:
+            datetime.strptime(end_date, "%Y%m%d")
+        except ValueError as err:
+            raise ValueError("Not a valid end date, use format 'YYYYMMDD'.") from err
+
+    # Check version make sure to include 'latest'
+    if version is not None:
+        if version != "latest" and (version[0] != "v" and len(version) != 4):
+            raise ValueError("Not a valid version, use format 'vXXX'.")
+
+    # check repointing follows 'repoint00000' format
+    if repointing is not None:
+        if not re.fullmatch(r"repoint\d{5}", str(repointing)):
+            raise ValueError(
+                "Not a valid repointing, use format repointing<num>,"
+                " where <num> is a 5 digit integer."
+            )
+
+    # check extension
+    if extension is not None:
+        if extension not in ("pkts", "cdf"):
+            raise ValueError("Not a valid extension, choose from ('pkts', 'cdf').")
+
     url = f"{imap_data_access.config['DATA_ACCESS_URL']}"
     url += f"/query?{urlencode(query_params)}"
 

tests/test_io.py

@@ -4,6 +4,7 @@
 
 import json
 import os
+import re
 import unittest
 from io import BytesIO
 from pathlib import Path
@@ -158,8 +159,8 @@ def test_download_already_exists(mock_urlopen: unittest.mock.MagicMock):
             "descriptor": "test-description",
             "start_date": "20100101",
             "end_date": "20100102",
-            "repointing": "00001",
-            "version": "000",
+            "repointing": "repoint00001",
+            "version": "v000",
             "extension": "pkts",
         },
         # Make sure not all query params are sent if they are missing
@@ -218,6 +219,67 @@ def test_query_bad_params(mock_urlopen: unittest.mock.MagicMock):
     assert mock_urlopen.call_count == 0
 
 
+@pytest.mark.parametrize(
+    ("query_flag", "query_input", "expected_output"),
+    [
+        # All parameters should  not send query
+        (
+            "instrument",
+            "badInput",
+            re.escape(
+                "Not a valid instrument, please choose from "
+                "('codice', 'glows', 'hi', 'hit', 'idex', 'lo',"
+                " 'mag', 'swapi', 'swe', 'ultra')"
+            ),
+        ),
+        (
+            "data_level",
+            "badInput",
+            re.escape(
+                "Not a valid data level, choose from "
+                "('l0','l1','l1a','l1b','l1c','l1ca','l1cb','l1d',"
+                "'l2','l2pre','l3','l3a','l3b','l3c','l3d')."
+            ),
+        ),
+        ("start_date", "badInput", "Not a valid start date, use format 'YYYYMMDD'."),
+        ("end_date", "badInput", "Not a valid end date, use format 'YYYYMMDD'."),
+        (
+            "repointing",
+            "badInput",
+            "Not a valid repointing, use format repointing<num>, "
+            "where <num> is a 5 digit integer.",
+        ),
+        ("version", "badInput", "Not a valid version, use format 'vXXX'."),
+        (
+            "extension",
+            "badInput",
+            re.escape("Not a valid extension, choose from ('pkts', 'cdf')."),
+        ),
+    ],
+)
+def test_bad_query_input(query_flag, query_input, expected_output):
+    """Test a function call to query with correct params but bad values.
+
+     Ensures correct error message is returned.
+
+    Parameters
+    ----------
+    query_flag : str
+        correct query flag.
+    query_input : str
+        incorrect query input.
+    expected_output : str
+        Output error expected to be given.
+    """
+    kwargs = {query_flag: query_input}
+
+    # Check if the ValueError is raised and contains the correct message
+    with pytest.raises(ValueError, match=expected_output):
+        imap_data_access.query(**kwargs)
+
+    # assert str(excinfo.value) == expected_output
+
+
 def test_upload_no_file(mock_urlopen: unittest.mock.MagicMock):
     """Test a call to the upload API that has no filename supplied.