From 81abeea27b266a9c56320ecab3ed2d5b34638ca6 Mon Sep 17 00:00:00 2001
From: "Ankur Sinha (Ankur Sinha Gmail)" <sanjay.ankur@gmail.com>
Date: Thu, 4 Jul 2024 09:45:45 +0100
Subject: [PATCH] feat(biosimulations): return URLs for convenience

---
 pyneuroml/biosimulations.py  | 39 +++++++++++++++++++++++++++---------
 tests/test_biosimulations.py |  3 ++-
 2 files changed, 31 insertions(+), 11 deletions(-)

diff --git a/pyneuroml/biosimulations.py b/pyneuroml/biosimulations.py
index c858b91e0..cab9cd1a3 100644
--- a/pyneuroml/biosimulations.py
+++ b/pyneuroml/biosimulations.py
@@ -168,7 +168,7 @@ def submit_simulation_archive(
     archive_file: str,
     sim_dict: typing.Dict[str, typing.Union[int, str, typing.List[str]]] = {},
     dry_run: bool = False,
-) -> object:
+) -> typing.Dict[str, typing.Union[str, object]]:
     """Submit an OMEX archive to biosimulations using the provided simulation run dictionary
 
     .. versionadded:: 1.2.10
@@ -203,7 +203,12 @@ def submit_simulation_archive(
         https://api.biosimulations.org/#/Simulations/SimulationRunController_createRun
 
     :type sim_dict: dict
-    :returns: the requests.post response object, or True if dry_run
+    :returns: dictionary with keys "response", "download", "logs", "view" that
+        contain the response object and download, logs, and view URLs
+
+        If a dry run, response is True and the URLs are all None
+
+    :rtype: dict
 
     """
     api_url = f"{biosimulations_api_url}/runs"
@@ -230,6 +235,8 @@ def submit_simulation_archive(
         logger.info(f"multipart encoded data is {m}")
         logger.info(f"with content type: {m.content_type}")
 
+        resdict = {}
+
         if dry_run is False:
             logger.info("Submitting archive to biosimulations")
             response = requests.post(
@@ -242,17 +249,29 @@ def submit_simulation_archive(
                 print(
                     f"Submitted {archive_file} successfully with id: {serv_response['id']}"
                 )
-                print(f"View: {biosimulations_api_url}/runs/{serv_response['id']}")
-                download_url = f'{biosimulations_api_url}/results/{serv_response["id"]}/download'
-                print(
-                    f"Downloads: {download_url}"
-                )
-                print(
-                    f"Logs: {biosimulations_api_url}/logs/{serv_response['id']}?includeOutput=true"
+
+                log_url = f"{biosimulations_api_url}/logs/{serv_response['id']}?includeOutput=true"
+                view_url = f"{biosimulations_api_url}/runs/{serv_response['id']}"
+                download_url = (
+                    f'{biosimulations_api_url}/results/{serv_response["id"]}/download'
                 )
+
+                print(f"View: {view_url}")
+                print(f"Downloads: {download_url}")
+                print(f"Logs: {log_url}")
+
+            resdict["response"] = response
+            resdict["view"] = view_url
+            resdict["download"] = download_url
+            resdict["logs"] = log_url
         else:
             response = True
             print("Dry run, not submitting")
             print(f"Simulation dictionary: {sim_dict}")
 
-    return download_url, response
+            resdict["response"] = response
+            resdict["view"] = None
+            resdict["download"] = None
+            resdict["logs"] = None
+
+    return resdict
diff --git a/tests/test_biosimulations.py b/tests/test_biosimulations.py
index 569f914f2..978104c69 100644
--- a/tests/test_biosimulations.py
+++ b/tests/test_biosimulations.py
@@ -62,9 +62,10 @@ def test_submit_simulation(self):
             "maxTime": "20",
         }
 
-        response = submit_simulation(
+        resdict = submit_simulation(
             "LEMS_NML2_Ex5_DetCell.xml", sim_dict=sim_dict, dry_run=dry_run
         )
+        response = resdict["response"]
         os.chdir(cwd)
 
         if dry_run: