From 662c35ba5ab0b944a5d13b4d3bd317014bff26ba Mon Sep 17 00:00:00 2001
From: Adam Reeve <adreeve@gmail.com>
Date: Sun, 18 Feb 2024 21:51:16 +1300
Subject: [PATCH] Add test case for reproducing the index error with truncated
 data

---
 nptdms/test/test_tdms_file.py | 30 ++++++++++++++++++++++++++++++
 nptdms/test/util.py           |  7 +++++--
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/nptdms/test/test_tdms_file.py b/nptdms/test/test_tdms_file.py
index 956a430..8be099b 100644
--- a/nptdms/test/test_tdms_file.py
+++ b/nptdms/test/test_tdms_file.py
@@ -764,6 +764,36 @@ def test_incomplete_segment_with_string_data():
     assert len(channel) == 0
 
 
+def test_truncated_interleaved_data():
+    """
+    Test when a segment is truncated within a row of interleaved data,
+    and the next segment offset is set but is beyond the end of the file.
+    """
+    test_file = GeneratedFile()
+    test_file.add_segment(
+        ("kTocMetaData", "kTocRawData", "kTocNewObjList", "kTocInterleavedData"),
+        segment_objects_metadata(
+            channel_metadata("/'group'/'channel1'", 3, 4),
+            channel_metadata("/'group'/'channel2'", 3, 4),
+        ),
+        "01 00 00 00" "02 00 00 00"
+        "03 00 00 00" "04 00 00 00"
+        "05 00 00 00" "06 00 00 00"
+        "07 00 00 00",
+        data_size_override=4 * 2 * 4
+    )
+    with test_file.get_tempfile() as temp_file:
+        with TdmsFile.open(temp_file.file) as tdms_file:
+            group = tdms_file['group']
+            chan1 = group['channel1']
+            chan2 = group['channel2']
+            for chan in [chan1, chan2]:
+                chan_data = chan[:]
+                assert chan[-1] == chan_data[-1]
+                assert len(chan) == 3
+                assert len(chan_data) == 3
+
+
 def test_truncated_metadata_in_last_segment():
     """ Test the scenario where writing the file was aborted with part of the metadata written
     """
diff --git a/nptdms/test/util.py b/nptdms/test/util.py
index a2c4b6d..7703965 100644
--- a/nptdms/test/util.py
+++ b/nptdms/test/util.py
@@ -223,7 +223,9 @@ class GeneratedFile(object):
     def __init__(self):
         self._content = []
 
-    def add_segment(self, toc, metadata, data, incomplete=False, binary_data=False, version=4713):
+    def add_segment(
+            self, toc, metadata, data, incomplete=False, binary_data=False, version=4713,
+            data_size_override=None):
         metadata_bytes = _hex_to_bytes(metadata)
         data_bytes = data if binary_data else _hex_to_bytes(data)
         if toc is not None:
@@ -246,7 +248,8 @@ def add_segment(self, toc, metadata, data, incomplete=False, binary_data=False,
                     raise ValueError("Unrecognised TOC value: %s" % toc_item)
             lead_in += struct.pack('<i', toc_mask)
             lead_in += struct.pack('<l', version)
-            next_segment_offset = len(metadata_bytes) + len(data_bytes)
+            data_len = data_size_override if data_size_override is not None else len(data_bytes)
+            next_segment_offset = len(metadata_bytes) + data_len
             raw_data_offset = len(metadata_bytes)
             if incomplete:
                 lead_in += _hex_to_bytes('FF' * 8)