From ebbfc8529da875200eab554a79817177d9d2ea63 Mon Sep 17 00:00:00 2001
From: Jermiah Joseph <jermiahjoseph98@gmail.com>
Date: Mon, 1 Apr 2024 12:22:42 -0400
Subject: [PATCH] feat: add refseries query to nbia client

---
 src/nbiatoolkit/nbia.py | 22 ++++++++++++++++++++++
 tests/test_tags.py      |  8 ++++++++
 2 files changed, 30 insertions(+)

diff --git a/src/nbiatoolkit/nbia.py b/src/nbiatoolkit/nbia.py
index 7da54ed..2225ac4 100644
--- a/src/nbiatoolkit/nbia.py
+++ b/src/nbiatoolkit/nbia.py
@@ -21,6 +21,13 @@
     ReturnType,
     conv_response_list,
 )
+
+from .dicomtags.tags import (
+    getReferencedSeriesUIDS,
+    extract_ROI_info,
+    getSequenceElement,
+)
+
 import pandas as pd
 import requests
 from requests.exceptions import JSONDecodeError as JSONDecodeError
@@ -615,6 +622,21 @@ def getDICOMTags(
 
         return conv_response_list(response, returnType)
 
+    def getRefSeriesUIDs(
+        self,
+        SeriesInstanceUID: str,
+    ) -> List[str]:
+
+        tags_df = self.getDICOMTags(
+            SeriesInstanceUID=SeriesInstanceUID,
+            return_type=ReturnType.DATAFRAME,
+        )
+
+        if type(tags_df) != pd.DataFrame:
+            raise ValueError("DICOM Tags not df or not found in the response.")
+
+        return getReferencedSeriesUIDS(series_tags_df=tags_df)
+
     def downloadSeries(
         self,
         SeriesInstanceUID: Union[str, list],
diff --git a/tests/test_tags.py b/tests/test_tags.py
index c3068bb..cdf350d 100644
--- a/tests/test_tags.py
+++ b/tests/test_tags.py
@@ -183,3 +183,11 @@ def test_extract_ROI_info(RTSTRUCT_Tags):
     ), "Expected all keys to be between 1 and 28"
 
     print("All test cases passed!")
+
+
+def test_getReferencedSeriesUIDS(client, RTSTRUCT_Series):
+    result = client.getRefSeriesUIDs(RTSTRUCT_Series["SeriesInstanceUID"].values[0])
+
+    expected = ["1.3.6.1.4.1.14519.5.2.1.133742245714270925254982946723351496764"]
+
+    assert result == expected, f"Expected {expected}, but got {result}"