Add test for reading BytesIO file

nptdms/test/tdms_test.py

@@ -1,14 +1,15 @@
 """Test reading of example TDMS files"""
 
-import unittest
-import sys
-import logging
 import binascii
-import struct
-import tempfile
 from datetime import datetime
-import os
+from io import BytesIO
+import logging
 import numpy as np
+import os
+import struct
+import sys
+import tempfile
+import unittest
 
 from nptdms import tdms
 
@@ -179,6 +180,12 @@ def load(self, *args, **kwargs):
         return tdms.TdmsFile(self.file, *args, **kwargs)
 
 
+class BytesIoTestFile(TestFile):
+    def __init__(self):
+        self.file = BytesIO()
+        self.data = bytes()
+
+
 class TDMSTestClass(unittest.TestCase):
     def setUp(self):
         logging.getLogger(tdms.__name__).setLevel(logging.DEBUG)
@@ -933,6 +940,22 @@ def test_raw_format(self):
                                         0.02380444, 0.20661031, 0.20447401,
                                         0.2517777])
 
+    def test_data_read_from_bytes_io(self):
+        """Test reading data"""
+
+        test_file = BytesIoTestFile()
+        test_file.add_segment(*basic_segment())
+        tdmsData = test_file.load()
+
+        data = tdmsData.channel_data("Group", "Channel1")
+        self.assertEqual(len(data), 2)
+        self.assertEqual(data[0], 1)
+        self.assertEqual(data[1], 2)
+        data = tdmsData.channel_data("Group", "Channel2")
+        self.assertEqual(len(data), 2)
+        self.assertEqual(data[0], 3)
+        self.assertEqual(data[1], 4)
+
 
 if __name__ == '__main__':
     unittest.main()