add-keep_until

1 files changed, 24 insertions, 1 deletions

diff --git a/src/argaze.test/DataStructures.py b/src/argaze.test/DataStructures.py
index bdb166b..ca46d5e 100644
--- a/src/argaze.test/DataStructures.py
+++ b/src/argaze.test/DataStructures.py
@@ -202,6 +202,27 @@ class TestTimeStampedBufferClass(unittest.TestCase):
         self.assertEqual(len(ts_buffer), 2)
         self.assertEqual(ts_buffer.first, (2, "C"))
 
+    def test_keep_until(self):
+        """Test TimeStampedBuffer keep_until method."""
+
+        ts_buffer = DataStructures.TimeStampedBuffer({0: "A", 1: "B", 2: "C", 3: "D"})
+
+        # Check keep_until an existing timestamp
+        keep_until_2 = ts_buffer.keep_until(2)
+        
+        self.assertEqual(keep_until_2, (2, "C"))
+        self.assertEqual(len(ts_buffer), 3)
+        self.assertEqual(ts_buffer.last, (2, "C"))
+
+        # Check keep_until an none existing timestamp
+        ts_buffer = DataStructures.TimeStampedBuffer({0: "A", 1: "B", 2: "C", 3: "D"})
+
+        keep_until_1dot5 = ts_buffer.keep_until(1.5)
+
+        self.assertEqual(keep_until_1dot5, (2, "C"))
+        self.assertEqual(len(ts_buffer), 3)
+        self.assertEqual(ts_buffer.last, (2, "C"))
+
     def test_last(self):
         """Test TimeStampedBuffer last property."""
 
@@ -266,7 +287,8 @@ class TestTimeStampedBufferClass(unittest.TestCase):
         with self.assertRaises(KeyError):
 
             ts_buffer.get_last_before(-1)
-
+            
+            
     def test_get_last_until(self):
         """Test TimeStampedBuffer get_last_until method."""
 
@@ -343,6 +365,7 @@ class TestTimeStampedBufferClass(unittest.TestCase):
         self.assertEqual(ts_buffer_dataframe.columns.size, 1)
         self.assertEqual(ts_buffer_dataframe.columns[0], "value")
 
+
 if __name__ == '__main__':
 
     unittest.main()
\ No newline at end of file