From 018799681525b9bf67a023e2793f165efc56cf51 Mon Sep 17 00:00:00 2001
From: Ruben Verweij <verweij@physics.leidenuniv.nl>
Date: Wed, 19 Jul 2017 11:29:55 +0200
Subject: [PATCH] Extend unit tests

---
 nd2reader/artificial.py | 30 +++++++++++++++++++++++++-----
 nd2reader/parser.py     |  5 +++--
 tests/test_reader.py    | 23 ++++++++++++++++++++---
 3 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/nd2reader/artificial.py b/nd2reader/artificial.py
index 3700c55..53d2677 100644
--- a/nd2reader/artificial.py
+++ b/nd2reader/artificial.py
@@ -22,12 +22,12 @@ class ArtificialND2(object):
                   'metadata_item': 11,
                   }
 
-    def __init__(self, file, version=(3, 0)):
+    def __init__(self, file, version=(3, 0), skip_blocks=None):
         self.version = version
         self.raw_text, self.locations, self.data = b'', None, None
         check_or_make_dir(path.dirname(file))
         self._fh = open(file, 'w+b', 0)
-        self.write_file()
+        self.write_file(skip_blocks)
 
     def __enter__(self):
         return self
@@ -51,9 +51,23 @@ class ArtificialND2(object):
         if self._fh is not None:
             self._fh.close()
 
-    def write_file(self):
-        self.write_version()
-        self.locations, self.data = self.write_label_map()
+    def write_file(self, skip_blocks=None):
+        if skip_blocks is None:
+            skip_blocks = []
+
+        if 'version' not in skip_blocks:
+            # write version header at start of the file
+            self.write_version()
+
+        if 'label_map' not in skip_blocks:
+            # write label map + data in the center
+            self.locations, self.data = self.write_label_map()
+
+        if 'label_map_marker' not in skip_blocks:
+            # write start position of label map at the end of the file
+            self.write_label_map_info()
+
+        # write all to file
         self._fh.write(self.raw_text)
 
     def write_version(self):
@@ -65,6 +79,12 @@ class ArtificialND2(object):
         # write version info
         self.raw_text += self._get_version_string()
 
+    def write_label_map_info(self):
+        """Write the location of the start of the label map at the end of the file
+        """
+        location = self._get_version_byte_length()
+        self.raw_text += struct.pack("Q", location)
+
     def _get_version_string(self):
         return six.b('ND2 FILE SIGNATURE CHUNK NAME01!Ver%s.%s' % self.version)
 
diff --git a/nd2reader/parser.py b/nd2reader/parser.py
index 6a98152..74a72bc 100644
--- a/nd2reader/parser.py
+++ b/nd2reader/parser.py
@@ -122,8 +122,8 @@ class Parser(object):
 
         """
         major_version, minor_version = get_version(self._fh)
-        supported = self.supported_file_versions.get((major_version, minor_version)) or \
-                    self.supported_file_versions.get((major_version, None))
+        supported = self.supported_file_versions.get(
+            (major_version, minor_version)) or self.supported_file_versions.get((major_version, None))
 
         if not supported:
             raise InvalidVersionError("No parser is available for that version.")
@@ -149,6 +149,7 @@ class Parser(object):
             LabelMap: the computed label map
 
         """
+        # go 8 bytes back from file end
         self._fh.seek(-8, 2)
         chunk_map_start_location = struct.unpack("Q", self._fh.read(8))[0]
         self._fh.seek(chunk_map_start_location)
diff --git a/tests/test_reader.py b/tests/test_reader.py
index 71b4033..ffadaae 100644
--- a/tests/test_reader.py
+++ b/tests/test_reader.py
@@ -1,14 +1,31 @@
 import unittest
-from os import path
 import numpy as np
 
 from nd2reader.artificial import ArtificialND2
-from nd2reader.common import check_or_make_dir
 from nd2reader.exceptions import EmptyFileError
-from nd2reader.parser import Parser
 from nd2reader.reader import ND2Reader
 
 
 class TestReader(unittest.TestCase):
     def test_extension(self):
         self.assertTrue('nd2' in ND2Reader.class_exts())
+
+    def test_init_and_init_axes(self):
+        with ArtificialND2('test_data/test_nd2_reader.nd2') as artificial:
+            reader = ND2Reader('test_data/test_nd2_reader.nd2')
+
+            attributes = artificial.data['image_attributes']['SLxImageAttributes']
+            self.assertEqual(reader.metadata['width'], attributes['uiWidth'])
+            self.assertEqual(reader.metadata['height'], attributes['uiHeight'])
+
+            self.assertEqual(reader.metadata['width'], reader.sizes['x'])
+            self.assertEqual(reader.metadata['height'], reader.sizes['y'])
+
+            self.assertEqual(reader._dtype, np.float64)
+            self.assertEqual(reader.iter_axes, ['t'])
+
+    def test_init_empty_file(self):
+        with ArtificialND2('test_data/empty.nd2', skip_blocks=['label_map_marker']):
+            with self.assertRaises(EmptyFileError) as exception:
+                ND2Reader('test_data/empty.nd2')
+            self.assertEqual(str(exception.exception), "No axes were found for this .nd2 file.")