From 6c971cf78ab7b099e93a601aade2d07626afbaa9 Mon Sep 17 00:00:00 2001
From: RemDelaporteMathurin <rdelaportemathurin@gmail.com>
Date: Tue, 17 Dec 2024 16:46:24 -0500
Subject: [PATCH] added an attribute to LSCReader

---
 libra_toolbox/tritium/lsc_measurements.py |  7 ++++++-
 test/tritium/test_lsc_reader.py           | 10 ++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/libra_toolbox/tritium/lsc_measurements.py b/libra_toolbox/tritium/lsc_measurements.py
index 60ac26a..e8f6845 100644
--- a/libra_toolbox/tritium/lsc_measurements.py
+++ b/libra_toolbox/tritium/lsc_measurements.py
@@ -11,6 +11,8 @@
 
 
 class LSCFileReader:
+    quench_set: str
+
     def __init__(
         self,
         file_path: str,
@@ -29,6 +31,7 @@ def __init__(
         self.labels_column = labels_column
         self.data = None
         self.header_content = None
+        self.quench_set = None
 
     def read_file(self):
         """
@@ -50,12 +53,14 @@ def read_file(self):
         with open(self.file_path, "r") as file:
             lines = file.readlines()
             for i, line in enumerate(lines):
+                if line.startswith("Quench Set:"):
+                    quench_set_line_idx = i + 1
                 if line.startswith("S#"):
                     start = i
                     break
                 header_lines.append(line)
         self.header_content = "".join(header_lines)
-
+        self.quench_set = lines[quench_set_line_idx].strip()
         # read the file with dataframe starting from the line with S#
         self.data = pd.read_csv(self.file_path, skiprows=start)
 
diff --git a/test/tritium/test_lsc_reader.py b/test/tritium/test_lsc_reader.py
index bd2f0e2..02ad3c1 100644
--- a/test/tritium/test_lsc_reader.py
+++ b/test/tritium/test_lsc_reader.py
@@ -47,3 +47,13 @@ def test_lsc_reader_label_automated():
 
     assert len(bq1_values) == 10
     assert len(bq1_values_with_labels) == 10
+
+
+def test_read_quench_set():
+    """Tests reading the quench set from the LSC file"""
+    filename = Path(__file__).parent / "test_lsc_file_with_labels.csv"
+
+    csv_reader = LSCFileReader(filename, labels_column="SMPL_ID")
+    csv_reader.read_file()
+
+    assert csv_reader.quench_set == "Low Energy: 3H-UG"