added functions, created tests

driver.py

@@ -0,0 +1,14 @@
+from src.nbia import NBIAClient
+import requests
+from pprint import pprint
+
+client = NBIAClient(log_level="INFO")
+
+response = client.getPatients(collection="LIDC-IDRI", modality="CT")
+
+pprint(response[0:5])
+
+
+response = client.getPatients(collection="LIDC-IDRI", modality="CR")
+
+pprint(response[0:5])

src/nbia.py

@@ -4,10 +4,14 @@
 import requests
 
 class NBIAClient:
-    def __init__(self, username: str = "nbia_guest", password: str = "") -> None:
+    def __init__(self, 
+                 username: str = "nbia_guest", 
+                 password: str = "",
+                 log_level: str = "INFO"
+                 ) -> None:
 
         # Setup logger
-        self.logger = setup_logger(name = "NBIAClient", console_logging=True, log_level='INFO')
+        self.logger = setup_logger(name = "NBIAClient", console_logging=True, log_level=log_level)
 
         # Setup OAuth2 client
         self.logger.info("Setting up OAuth2 client...")
@@ -30,21 +34,58 @@ def query_api(self, endpoint: NBIA_ENDPOINTS, params: dict = {}) -> dict:
         query_url = base_url + endpoint.value
 
         self.logger.info("Querying API endpoint: %s", query_url)
-        response = requests.get(query_url, headers=self.api_headers).json()
-
+        self.logger.debug("API headers: %s", (self._createDebugURL(endpoint, params)))
+        response = requests.get(
+            url=query_url, 
+            headers=self.api_headers,
+            params=params
+            ).json()
+
+        self.logger.debug("API response: %s", response)
         return response
-
-    def getCollectionPatientCount(self) -> dict:
-        return self.query_api(NBIA_ENDPOINTS.GET_COLLECTION_PATIENT_COUNT)
 
-    def getCollections(self) -> dict:
-        return self.query_api(NBIA_ENDPOINTS.GET_COLLECTIONS)
+    def _createDebugURL(self, endpoint, params):
+        auth = "'Authorization:" + self.api_headers["Authorization"] + "' -k "
+        base_url = "'https://services.cancerimagingarchive.net/nbia-api/services/"
+        query_url = auth + base_url + endpoint.value
+        debug_url = query_url + "?"
+        for key, value in params.items():
+            debug_url += f"{key}={value}&"
+        debug_url = debug_url[:-1] + "'"
+        return debug_url
 
-    def getBodyPartCounts(self, collection: str = "", modality: str = "") -> dict:
+    def getCollections(self) -> list:
+        response = self.query_api(NBIA_ENDPOINTS.GET_COLLECTIONS)
+        collections = []
+        for collection in response:
+            collections.append(collection["Collection"])
+        return collections
+
+    def getCollectionPatientCount(self) -> list:
+        response = self.query_api(NBIA_ENDPOINTS.GET_COLLECTION_PATIENT_COUNT)
+        patientCount = []
+        for collection in response:
+            patientCount.append({"Collection": collection["criteria"], "PatientCount": collection["count"]})
+        return patientCount
+
+    def getBodyPartCounts(self, collection: str = "", modality: str = "") -> list:
         PARAMS = {}
         if collection:
             PARAMS["Collection"] = collection
         if modality:
             PARAMS["Modality"] = modality
+        response =  self.query_api(endpoint = NBIA_ENDPOINTS.GET_BODY_PART_PATIENT_COUNT, params = PARAMS)
+        bodyparts=[]
+        for bodypart in response:
+            bodyparts.append({"BodyPartExamined": bodypart["criteria"], "Count": bodypart["count"]})
+        return bodyparts
+
+    def getPatients(self, collection: str, modality: str) -> list:
+        assert collection is not None
+        assert modality is not None
 
-        return self.query_api(endpoint = NBIA_ENDPOINTS.GET_BODY_PART_PATIENT_COUNT, params = PARAMS)
+        PARAMS = {"Collection": collection,"Modality": modality}
+        response = self.query_api(endpoint = NBIA_ENDPOINTS.GET_PATIENT_BY_COLLECTION_AND_MODALITY, params = PARAMS)
+        patientList = [_["PatientId"] for _ in response]
+        return patientList
+

src/tests/oldtest_nbia.py

@@ -24,7 +24,7 @@ def test_nbiaclient_access_token(nbia_client):
 
 def test_nbia_getCollections(nbia_client_collections):
     collections = nbia_client_collections._collections
-    assert isinstance(collections, pd.DataFrame)
+    assert isinstance(collections, dict)
     assert collections.shape[0] > 0
     assert collections.shape[1] == 1
     assert collections.columns[0] == 'Collection'

src/tests/test_nbia.py

@@ -0,0 +1,57 @@
+#####
+# pytest -v -n 8
+import pytest
+from ..nbia import NBIAClient
+
+@pytest.fixture(scope="session")
+def nbia_client():
+    nbia = NBIAClient()
+    return nbia
+
+@pytest.fixture(scope="session")
+def nbia_client_bad_username():
+    nbia = NBIAClient(username="bad_username", password="bad_password")
+    return nbia
+
+def test_nbiaclient_access_token(nbia_client):
+    assert nbia_client.api_headers is not None
+
+def test_getCollections(nbia_client):
+    collections = nbia_client.getCollections()
+    assert isinstance(collections, list)
+    assert len(collections) > 0
+
+def test_getBodyPartCounts_all(nbia_client):
+    bodyparts = nbia_client.getBodyPartCounts()
+    assert isinstance(bodyparts, list)
+    assert len(bodyparts) > 0
+    assert "BodyPartExamined" in bodyparts[0]    
+    assert "Count" in bodyparts[0]
+    assert int(bodyparts[0]["Count"]) > 0
+
+def test_getBodyPartCounts_4DLung(nbia_client):
+    bodyparts = nbia_client.getBodyPartCounts(collection="4D-Lung")
+    assert isinstance(bodyparts, list)
+    assert len(bodyparts) > 0
+    assert isinstance(bodyparts[0], dict)
+    assert "BodyPartExamined" in bodyparts[0]    
+    assert "Count" in bodyparts[0]
+    assert bodyparts[0]["BodyPartExamined"] == "LUNG"
+    assert int(bodyparts[0]["Count"]) > 0
+
+def test_getCollectionPatientCount(nbia_client):
+    patientCount = nbia_client.getCollectionPatientCount()
+    assert isinstance(patientCount, list)
+    assert len(patientCount) > 0
+    assert isinstance(patientCount[0], dict)
+    assert "Collection" in patientCount[0]
+    assert "PatientCount" in patientCount[0]
+    assert int(patientCount[0]["PatientCount"]) > 0
+
+def test_getPatients(nbia_client):
+    patientList = nbia_client.getPatients("LIDC-IDRI", "CT")
+    assert isinstance(patientList, list)
+    assert len(patientList) > 0
+    assert isinstance(patientList[0], str)
+    assert len(patientList[0]) > 0
+

src/utils/nbia_endpoints.py

@@ -9,7 +9,7 @@ class NBIA_ENDPOINTS(Enum):
     GET_COLLECTION_PATIENT_COUNT = 'getCollectionValuesAndCounts'
     GET_COLLECTIONS = 'v2/getCollectionValues'
     GET_BODY_PART_PATIENT_COUNT = 'getBodyPartValuesAndCounts'
-    GET_PATIENT_BY_COLLECTION_AND_MODALITY = 'getPatientByCollectionAndModality'
+    GET_PATIENT_BY_COLLECTION_AND_MODALITY = 'v2/getPatientByCollectionAndModality'
 
 
     # Helper functions