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 @@
         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 @@
     # 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 @@
         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 @@
         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 @@
         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 @@
         """
 
         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 @@
                     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 @@
             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 @@
 
     :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 @@
         access_keys=[],
         test=False,
         interactive_mode=True,
+        refresh=False,
     ):
         """
         Manages credentials for OLIApi authentication requests.
@@ -102,6 +103,7 @@
         :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
@@ -121,6 +123,8 @@
             _logger.setLevel(logging.DEBUG)
         self.set_headers()
 
+        self.refresh = refresh
+
     def set_headers(self):
         """
         Creates OLI Cloud API headers and performs initial login.
@@ -191,7 +195,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):
@@ -336,12 +340,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,7 +358,7 @@
             )
         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 @@
                     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.")

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(