Repair OLIAPI

tutorials/oli_calculations.ipynb

@@ -273,7 +273,7 @@
     "credential_manager = CredentialManager(\n",
     "    access_keys=[os.environ[\"OLI_API_KEY\"]],\n",
     "    root_url=os.environ[\"OLI_API_ROOT_URL\"],\n",
-    "    interactive_mode=True,\n",
+    "    interactive_mode=False,\n",
     "    test=True,\n",
     ")\n",
     "'''\n",
@@ -411,16 +411,18 @@
     "    dbs_files = oliapi.get_user_dbs_file_ids()\n",
     "    \n",
     "    # uncomment the following lines to view results:\n",
-    "    #for idx, file in enumerate(dbs_files):\n",
-    "    #    print(f\"{idx+1}\\t{file}\")\n",
+    "    for idx, file in enumerate(dbs_files):\n",
+    "       print(f\"{idx+1}\\t{file}\")\n",
     "    \n",
     "    # if a DBS file has been retained from a previous session,\n",
     "    # its flash history and chemistry information can be summarized.\n",
     "    file_summary = oliapi.get_dbs_file_summary(dbs_file_id)\n",
-    "\n",
+    "    print(file_summary)\n",
+    "    \n",
     "    # save chemistry information\n",
     "    chemistry_info = file_summary[\"chemistry_info\"]\n",
-    "    write_output(chemistry_info[\"result\"], \"chemistry_info.json\")\n",
+    "    # Uncomment next line to save as json file\n",
+    "    # write_output(chemistry_info[\"result\"], \"chemistry_info.json\")\n",
     "    \n",
     "    # DBS files can also be deleted.\n",
     "    oliapi.dbs_file_cleanup(dbs_files)\n",

watertap/tools/oli_api/client.py

@@ -52,10 +52,10 @@
 import sys
 import json
 import time
-
 from pyomo.common.dependencies import attempt_import
 
 requests, requests_available = attempt_import("requests", defer_check=False)
+from requests import JSONDecodeError
 
 from watertap.tools.oli_api.util.watertap_to_oli_helper_functions import get_oli_name
 
@@ -95,6 +95,9 @@
         else:
             _logger.setLevel(logging.DEBUG)
 
+        # TODO: unknown bug where only "liquid1" phase is found in Flash analysis
+        self.valid_phases = ["liquid1", "vapor", "solid", "liquid2"]
+
     # binds OLIApi instance to context manager
     def __enter__(self):
         self.session_dbs_files = []
@@ -103,13 +106,18 @@
     # return False if no exceptions raised
     def __exit__(self, exc_type=None, exc_value=None, traceback=None):
         # delete all .dbs files created during session
+        _logger.info(
+            f"Exiting: deleting {len(self.session_dbs_files)} remaining DBS files created during the session that were not marked by keep_file=True."
+        )
         self.dbs_file_cleanup(self.session_dbs_files)
         return False
 
     def _prompt(self, msg, default=""):
         if self.interactive_mode:
+            msg = msg + "Enter [y]/n to proceed."
             return input(msg)
         else:
+            msg = msg + "To choose [y]/n to this action, set interactive_mode=True."
             _logger.info(msg)
             return default
 
@@ -133,7 +141,7 @@
         if bool(dbs_file_id):
             if not keep_file:
                 self.session_dbs_files.append(dbs_file_id)
-            _logger.info(f"DBS file ID is {dbs_file_id}")
+            _logger.info(f"Uploaded DBS file ID is {dbs_file_id}")
             return dbs_file_id
         else:
             raise RuntimeError("Unexpected failure getting DBS file ID.")
@@ -191,8 +199,6 @@
         else:
             dbs_file_inputs["modelName"] = "OLI_analysis"
 
-        # TODO: unknown bug where only "liquid1" phase is found in Flash analysis
-        self.valid_phases = ["liquid1", "vapor", "solid", "liquid2"]
         if phases is not None:
             invalid_phases = [p for p in phases if p not in self.valid_phases]
             if invalid_phases:
@@ -230,7 +236,7 @@
         if bool(dbs_file_id):
             if not keep_file:
                 self.session_dbs_files.append(dbs_file_id)
-            _logger.info(f"DBS file ID is {dbs_file_id}")
+            _logger.info(f"Generated DBS file ID is {dbs_file_id}")
             return dbs_file_id
         else:
             raise RuntimeError("Unexpected failure getting DBS file ID.")
@@ -289,10 +295,15 @@
         """
 
         if dbs_file_ids is None:
+            _logger.info(
+                "No DBS file IDs were provided to the dbs_file_cleanup method. Checking user's cloud account for DBS file IDs."
+            )
             dbs_file_ids = self.get_user_dbs_file_ids()
-        r = self._prompt(
-            f"WaterTAP will delete {len(dbs_file_ids)} DBS files [y]/n ", "y"
-        )
+            if not len(dbs_file_ids):
+                _logger.info("No DBS file IDs were found on the user's cloud account.")
+                return
+
+        r = self._prompt(f"WaterTAP will delete {len(dbs_file_ids)} DBS files. ", "y")
         if (r.lower() == "y") or (r == ""):
             for dbs_file_id in dbs_file_ids:
                 _logger.info(f"Deleting {dbs_file_id} ...")
@@ -302,8 +313,16 @@
                     headers=self.credential_manager.headers,
                 )
                 req = _request_status_test(req, ["SUCCESS"])
+
                 if req["status"] == "SUCCESS":
-                    _logger.info(f"File deleted")
+                    # Remove the file from session_dbs_files list if it is there, otherwise an error will occur upon exit when this method is called again and already deleted files will remain on the list for deletion. Thus, an error can occur if there is no existing ID to delete.
+                    if dbs_file_id in self.session_dbs_files:
+                        self.session_dbs_files.remove(dbs_file_id)
+                        _logger.info(
+                            f"File {dbs_file_id} deleted and removed from session_dbs_files list."
+                        )
+                    else:
+                        _logger.info(f"File {dbs_file_id} deleted.")
 
     def get_corrosion_contact_surfaces(self, dbs_file_id):
         """
@@ -428,10 +447,21 @@
         """
 
         mode, url, headers = self._get_flash_mode(dbs_file_id, flash_method)
-        req = requests.request(
-            mode, url, headers=headers, data=json.dumps(input_params)
-        )
-        req_json = _request_status_test(req, ["SUCCESS"])
+        try:
+            req = requests.request(
+                mode, url, headers=headers, data=json.dumps(input_params)
+            )
+            req_json = _request_status_test(req, ["SUCCESS"])
+        except JSONDecodeError:
+            delay = 1
+            time.sleep(delay)
+            _logger.debug(
+                f"JSONDecodeError occurred. Trying to convert response object to JSON again."
+            )
+            req = requests.request(
+                mode, url, headers=headers, data=json.dumps(input_params)
+            )
+            req_json = _request_status_test(req, ["SUCCESS"])
         result_link = _get_result_link(req_json)
         result = _poll_result_link(result_link, headers, max_request, poll_time)
         return result
@@ -454,7 +484,7 @@
             if "resultsLink" in req_json["data"]:
                 result_link = req_json["data"]["resultsLink"]
     if not result_link:
-        raise RuntimeError(f"Failed to get 'resultsLink'. Response: {req.json()}")
+        raise RuntimeError(f"Failed to get 'resultsLink'. Response: {req_json.json()}")
     return result_link
 
 
@@ -467,18 +497,21 @@
 
     :return req_json: response object converted to JSON
     """
-
     req_json = req.json()
+
     func_name = sys._getframe().f_back.f_code.co_name
     _logger.debug(f"{func_name} response: {req_json}")
+
     if req.status_code == 200:
         if target_keys:
             if "status" in req_json:
                 if req_json["status"] in target_keys:
                     return req_json
         else:
             return req_json
-    raise RuntimeError(f"Failure in {func_name}. Response: {req_json}")
+    raise RuntimeError(
+        f"Failure in {func_name}. Response: {req_json}. Status Code: {req.status_code}."
+    )
 
 
 def _poll_result_link(result_link, headers, max_request, poll_time):

watertap/tools/oli_api/conftest.py

@@ -82,8 +82,6 @@ def auth_credentials() -> dict:
 def oliapi_instance(
     tmp_path: Path,
     auth_credentials: dict,
-    local_dbs_file: Path,
-    source_water: dict,
 ) -> OLIApi:
 
     if not cryptography_available:
@@ -93,11 +91,11 @@ def oliapi_instance(
     credentials = {
         **auth_credentials,
         "config_file": cred_file_path,
+        "refresh": True,
+        "interactive_mode": False,
     }
     credential_manager = CredentialManager(**credentials, test=True)
     with OLIApi(credential_manager, interactive_mode=False) as oliapi:
-        oliapi.upload_dbs_file(str(local_dbs_file))
-        oliapi.generate_dbs_file(source_water)
         yield oliapi
     with contextlib.suppress(FileNotFoundError):
         cred_file_path.unlink()

watertap/tools/oli_api/credentials.py

@@ -72,7 +72,6 @@
 _logger.setLevel(logging.DEBUG)
 
 
-# TODO: consider letting user set logging level instead of using interactive mode
 class CredentialManager:
     """
     A class to handle credentials for OLI Cloud.
@@ -89,6 +88,8 @@
         access_keys=[],
         test=False,
         interactive_mode=True,
+        refresh=False,
+        debug_level="INFO",
     ):
         """
         Manages credentials for OLIApi authentication requests.
@@ -101,24 +102,30 @@
         :param encryption_key: fernet key generated by credential manager object
         :param access_keys: list of access keys generated by user
         :param test: bool switch for automation during tests
-        :param interactive_mode: bool to switch level of logging display from info to debug only
+        :param interactive_mode: enables direct interaction with user through prompts
+        :param refresh: bool to use refresh token in login method
+        :param debug_level: string defining level of logging activity
+
         """
 
         self.test = test
         self.access_key = ""
         self.encryption_key = encryption_key
         self.config_file = Path(config_file).resolve()
+        self.refresh = refresh
+        self.interactive_mode = interactive_mode
+        self.debug_level = debug_level
+        if self.debug_level == "INFO":
+            _logger.setLevel(logging.INFO)
+        else:
+            _logger.setLevel(logging.DEBUG)
         self._manage_credentials(
             username,
             password,
             root_url,
             auth_url,
             access_keys,
         )
-        if interactive_mode:
-            _logger.setLevel(logging.INFO)
-        else:
-            _logger.setLevel(logging.DEBUG)
         self.set_headers()
 
     def set_headers(self):
@@ -191,7 +198,7 @@
         self.dbs_url = self.credentials["root_url"] + "/channel/dbs"
         self.upload_dbs_url = self.credentials["root_url"] + "/channel/upload/dbs"
         self.access_key_url = self.credentials["root_url"] + "/user/api-key"
-        self.engine_url = self.credentials["root_url"] + "/engine/"
+        self.engine_url = self.credentials["root_url"] + "/engine"
         self._delete_dbs_url = self.credentials["root_url"] + "/channel/file/"
 
     def _decrypt_credentials(self):
@@ -233,12 +240,13 @@
         if self.test:
             return True
         else:
-            r = input(
-                "WaterTAP will write encrypted file to store OLI Cloud credentials: [y]/n: "
-            )
-            if (r.lower() == "y") or (r == ""):
-                return True
-            return False
+            if self.interactive_mode:
+                r = input(
+                    "WaterTAP will write encrypted file to store OLI Cloud credentials: [y]/n: "
+                )
+                if (r.lower() == "y") or (r == ""):
+                    return True
+                return False
 
     # TODO: consider updating credentials when writing, rather than resetting blank ones
     def _encrypt_credentials(self):
@@ -336,12 +344,10 @@
         _logger.info(response.text)
         return response.text
 
-    def login(self, refresh=False):
+    def login(self):
         """
         Log in to OLI Cloud using access key or credentials.
 
-        :param refresh: bool to get refresh token
-
         :return status: bool indicating success or failure
         """
 
@@ -356,12 +362,22 @@
             )
         else:
             _logger.info("Logging into OLI API using username and password")
-            if refresh:
-                body = {
-                    "refresh_token": self.refresh_token,
-                    "grant_type": "refresh_token",
-                    "client_id": "apiclient",
-                }
+            if self.refresh:
+                if hasattr(self, "refresh_token"):
+                    body = {
+                        "refresh_token": self.refresh_token,
+                        "grant_type": "refresh_token",
+                        "client_id": "apiclient",
+                    }
+                else:
+                    body = {
+                        "username": self.credentials["username"],
+                        "password": self.credentials["password"],
+                        "grant_type": "password",
+                        "client_id": "apiclient",
+                    }
+                    status = self.auth_status(body, req_result)
+                    return status
             else:
                 body = {
                     "username": self.credentials["username"],

watertap/tools/oli_api/flash.py

@@ -972,7 +972,9 @@
                     pass
                 elif k in d["inflows"]["values"]:
                     d = d["inflows"]["values"]
-                elif k in d["corrosionParameters"]:
+                elif hasattr(d, "corrosionParameters") and (
+                    k in d["corrosionParameters"]
+                ):
                     d = d["corrosionParameters"]
                 else:
                     _logger.warning(f"Survey key {k} not found in JSON input.")
@@ -1158,9 +1160,9 @@
                         prop_tag = _get_nested_data(result, prop)["name"]
             else:
                 _logger.warning(
-                    f"Unexpected result:\n{result}\n\ninput_dict:\n{input_dict}"
+                    f"Unexpected result:\n{result}\n\ninput_dict:\n{input_dict} from prop {prop}"
                 )
-
+                continue
             label = f"{prop_tag}_{phase_tag}" if phase_tag else prop_tag
             input_dict[k][label] = _extract_values(result, prop)
         return input_dict

watertap/tools/oli_api/tests/test_client.py

@@ -56,8 +56,26 @@ def test_dbs_file_available_for_testing(local_dbs_file: Path):
     assert local_dbs_file.is_file()
 
 
+@pytest.mark.unit
+def test_generate_dbs_file(
+    oliapi_instance: OLIApi, local_dbs_file: Path, source_water: dict
+):
+    dbs_file_id = oliapi_instance.generate_dbs_file(source_water)
+    assert len(dbs_file_id) > 0
+
+
+@pytest.mark.unit
+def test_upload_dbs_file(
+    oliapi_instance: OLIApi, local_dbs_file: Path, source_water: dict
+):
+    dbs_file_id = oliapi_instance.upload_dbs_file(str(local_dbs_file))
+    assert len(dbs_file_id) > 0
+
+
 @pytest.mark.unit
 def test_dbs_file_cleanup(oliapi_instance: OLIApi, local_dbs_file: Path):
+    # This test checks both the upload_dbs_file method and dbs_file_cleanup method
+    # The following line will return 3 DBS file IDs. Note, the same file is uploaded three times, but each time a new ID is assigned.
     ids = [oliapi_instance.upload_dbs_file(str(local_dbs_file)) for i in range(3)]
     oliapi_instance.dbs_file_cleanup(ids)