Repair OLIAPI

tutorials/oli_calculations.ipynb

@@ -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

@@ -56,7 +56,6 @@
 from pyomo.common.dependencies import attempt_import
 
 requests, requests_available = attempt_import("requests", defer_check=False)
-
 from watertap.tools.oli_api.util.watertap_to_oli_helper_functions import get_oli_name
 
 
@@ -95,6 +94,9 @@ def __init__(self, credential_manager, interactive_mode=True, debug_level="INFO"
         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 +105,18 @@ def __enter__(self):
     # 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 +140,7 @@ def upload_dbs_file(self, dbs_file_path, keep_file=False):
         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 +198,6 @@ def generate_dbs_file(
         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 +235,7 @@ def generate_dbs_file(
         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 +294,15 @@ def dbs_file_cleanup(self, dbs_file_ids=None):
         """
 
         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 +312,16 @@ def dbs_file_cleanup(self, dbs_file_ids=None):
                     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):
         """
@@ -454,7 +472,7 @@ def _get_result_link(req_json):
             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 +485,20 @@ def _request_status_test(req, target_keys):
 
     :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

@@ -89,6 +89,7 @@ def __init__(
         access_keys=[],
         test=False,
         interactive_mode=True,
+        refresh=False,
     ):
         """
         Manages credentials for OLIApi authentication requests.
@@ -102,12 +103,14 @@ def __init__(
         :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 refresh: bool to use refresh token in login method
         """
 
         self.test = test
         self.access_key = ""
         self.encryption_key = encryption_key
         self.config_file = Path(config_file).resolve()
+        self.refresh = refresh
         self._manage_credentials(
             username,
             password,
@@ -191,7 +194,7 @@ def _manage_credentials(self, username, password, root_url, auth_url, access_key
         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):
@@ -336,12 +339,10 @@ def delete_oliapi_access_key(self, api_key):
         _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,7 +357,7 @@ def login(self, refresh=False):
             )
         else:
             _logger.info("Logging into OLI API using username and password")
-            if refresh:
+            if self.refresh:
                 body = {
                     "refresh_token": self.refresh_token,
                     "grant_type": "refresh_token",

watertap/tools/oli_api/flash.py

@@ -972,7 +972,9 @@ def get_clone(self, flash_method, json_input, index, survey=None):
                     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 @@ def _create_input_dict(props, result):
                         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)
 

watertap/tools/oli_api/tests/test_flash.py

@@ -58,7 +58,7 @@
 def test_water_analysis_single_point(
     flash_instance: Flash, source_water: dict, oliapi_instance: OLIApi, tmp_path: Path
 ):
-    dbs_file_id = oliapi_instance.session_dbs_files[-1]
+    dbs_file_id = oliapi_instance.generate_dbs_file(source_water)
     stream_input = flash_instance.configure_water_analysis(
         source_water,
         file_name=tmp_path / "test_wa_input",
@@ -76,7 +76,7 @@ def test_water_analysis_single_point(
 def test_water_analysis_survey(
     flash_instance: Flash, source_water: dict, oliapi_instance: OLIApi, tmp_path: Path
 ):
-    dbs_file_id = oliapi_instance.session_dbs_files[-1]
+    dbs_file_id = oliapi_instance.generate_dbs_file(source_water)
     survey = build_survey(
         {
             "Na_+": linspace(0, 1e4, 2),
@@ -100,7 +100,7 @@ def test_water_analysis_survey(
 def test_isothermal_flash_single_point(
     flash_instance: Flash, source_water: dict, oliapi_instance: OLIApi, tmp_path: Path
 ):
-    dbs_file_id = oliapi_instance.session_dbs_files[-1]
+    dbs_file_id = oliapi_instance.generate_dbs_file(source_water)
     stream_input = flash_instance.configure_water_analysis(source_water)
     inflows = flash_instance.get_apparent_species_from_true(
         stream_input,
@@ -119,7 +119,7 @@ def test_isothermal_flash_single_point(
 def test_isothermal_flash_survey(
     flash_instance: Flash, source_water: dict, oliapi_instance: OLIApi, tmp_path: Path
 ):
-    dbs_file_id = oliapi_instance.session_dbs_files[-1]
+    dbs_file_id = oliapi_instance.generate_dbs_file(source_water)
     survey = build_survey(
         {
             "NaCl": linspace(0, 1e4, 2),
@@ -128,7 +128,6 @@ def test_isothermal_flash_survey(
         get_oli_names=True,
         file_name=tmp_path / "test_survey",
     )
-    dbs_file_id = oliapi_instance.session_dbs_files[-1]
     stream_input = flash_instance.configure_water_analysis(source_water)
     inflows = flash_instance.get_apparent_species_from_true(
         stream_input,
@@ -148,7 +147,7 @@ def test_isothermal_flash_survey(
 def test_bubble_point(
     flash_instance: Flash, source_water: dict, oliapi_instance: OLIApi, tmp_path: Path
 ):
-    dbs_file_id = oliapi_instance.session_dbs_files[-1]
+    dbs_file_id = oliapi_instance.generate_dbs_file(source_water)
 
     stream_input = flash_instance.configure_water_analysis(source_water)
     inflows = flash_instance.get_apparent_species_from_true(