From da38f7cd6673974cd46554e7936f4558785b7702 Mon Sep 17 00:00:00 2001
From: Ruben Verweij <verweij@physics.leidenuniv.nl>
Date: Thu, 30 Mar 2017 15:28:49 +0200
Subject: [PATCH] Add more unit test for common + fix possibly endless loop

---
 nd2reader/common.py      | 15 ++++++++++++---
 tests/test_artificial.py | 22 ++++++++++++++++++++++
 tests/test_common.py     | 29 ++++++++++++++++++++++++++---
 3 files changed, 60 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_artificial.py

diff --git a/nd2reader/common.py b/nd2reader/common.py
index c6ba459..fbb1c4a 100644
--- a/nd2reader/common.py
+++ b/nd2reader/common.py
@@ -155,10 +155,19 @@ def _parse_string(data):
 
         """
     value = data.read(2)
+    # the string ends at the first instance of \x00\x00
     while not value.endswith(six.b("\x00\x00")):
-        # the string ends at the first instance of \x00\x00
-        value += data.read(2)
-    return value.decode("utf16")[:-1].encode("utf8")
+        next_data = data.read(2)
+        if len(next_data) == 0:
+            break
+        value += next_data
+
+    try:
+        decoded = value.decode("utf16")[:-1].encode("utf8")
+    except UnicodeDecodeError:
+        decoded = value.decode('utf8')
+
+    return decoded
 
 
 def _parse_char_array(data):
diff --git a/tests/test_artificial.py b/tests/test_artificial.py
new file mode 100644
index 0000000..e9d2b29
--- /dev/null
+++ b/tests/test_artificial.py
@@ -0,0 +1,22 @@
+import unittest
+from os import path
+import six
+import struct
+
+from nd2reader.artificial import ArtificialND2
+from nd2reader.common import get_version, parse_version, parse_date, _add_to_metadata, _parse_unsigned_char, \
+    _parse_unsigned_int, _parse_unsigned_long, _parse_double, check_or_make_dir
+from nd2reader.exceptions import InvalidVersionError
+
+
+class TestArtificial(unittest.TestCase):
+    def setUp(self):
+        dir_path = path.dirname(path.realpath(__file__))
+        check_or_make_dir(path.join(dir_path, 'test_data/'))
+        self.test_file = path.join(dir_path, 'test_data/test.nd2')
+        self.create_test_nd2()
+
+    def create_test_nd2(self):
+        with ArtificialND2(self.test_file) as artificial:
+            self.assertIsNotNone(artificial.file_handle)
+            artificial.close()
diff --git a/tests/test_common.py b/tests/test_common.py
index 5dea7a0..f08ce7f 100644
--- a/tests/test_common.py
+++ b/tests/test_common.py
@@ -1,11 +1,14 @@
 import unittest
 from os import path
+
+import array
 import six
 import struct
 
 from nd2reader.artificial import ArtificialND2
 from nd2reader.common import get_version, parse_version, parse_date, _add_to_metadata, _parse_unsigned_char, \
-    _parse_unsigned_int, _parse_unsigned_long, _parse_double, check_or_make_dir
+    _parse_unsigned_int, _parse_unsigned_long, _parse_double, check_or_make_dir, _parse_string, _parse_char_array, \
+    get_from_dict_if_exists
 from nd2reader.exceptions import InvalidVersionError
 
 
@@ -78,9 +81,9 @@ class TestCommon(unittest.TestCase):
         self.assertDictEqual(metadata, {'test': ['value1', 'value2', 'value3']})
 
     @staticmethod
-    def _prepare_bin_stream(binary_format, value):
+    def _prepare_bin_stream(binary_format, *value):
         file = six.BytesIO()
-        data = struct.pack(binary_format, value)
+        data = struct.pack(binary_format, *value)
         file.write(data)
         file.seek(0)
         return file
@@ -97,3 +100,23 @@ class TestCommon(unittest.TestCase):
 
         file = self._prepare_bin_stream("d", 47.9)
         self.assertEqual(_parse_double(file), 47.9)
+
+        test_string = 'colloid'
+        file = self._prepare_bin_stream("%ds" % len(test_string), six.b(test_string))
+        parsed = _parse_string(file)
+        self.assertEqual(parsed, test_string)
+
+        test_data = [1, 2, 3, 4, 5]
+        file = self._prepare_bin_stream("Q" + ''.join(['B'] * len(test_data)), len(test_data), *test_data)
+        parsed = _parse_char_array(file)
+        self.assertEqual(parsed, array.array('B', test_data))
+
+    def test_get_from_dict_if_exists(self):
+        test_dict = {
+            six.b('existing'): 'test',
+            'string': 'test2'
+        }
+
+        self.assertIsNone(get_from_dict_if_exists('nowhere', test_dict))
+        self.assertEqual(get_from_dict_if_exists('existing', test_dict), 'test')
+        self.assertEqual(get_from_dict_if_exists('string', test_dict, convert_key_to_binary=False), 'test2')