Added tests for binary file creation

* fixed return value mismatch in get_xex to create_executable_file_image

atrcopy/ataridos.py

@@ -650,7 +650,8 @@ class AtariDosDiskImage(DiskImageBase):
         return segments_out
 
     def create_executable_file_image(self, segments, run_addr=None):
-        return get_xex(segments, run_addr), "XEX"
+        base_segment, user_segments = get_xex(segments, run_addr)
+        return base_segment, "XEX"
 
 
 class BootDiskImage(AtariDosDiskImage):
@@ -719,7 +720,7 @@ class BootDiskImage(AtariDosDiskImage):
         return []
 
 
-def get_xex(segments, runaddr=None):
+def get_xex(segments, run_addr=None):
     segments_copy = [s for s in segments]  # don't affect the original list!
     main_segment = None
     sub_segments = []
@@ -732,9 +733,17 @@ def get_xex(segments, runaddr=None):
             runad = True
     if not runad:
         words = np.empty([1], dtype='<u2')
-        if not runaddr:
-            runaddr = segments[0].start_addr
-        words[0] = runaddr
+        if run_addr:
+            found = False
+            for s in segments:
+                if run_addr >= s.start_addr and run_addr < s.start_addr + len(s):
+                    found = True
+                    break
+            if not found:
+                raise InvalidBinaryFile("Run address points outside data segments")
+        else:
+            run_addr = segments[0].start_addr
+        words[0] = run_addr
         r = SegmentData(words.view(dtype=np.uint8))
         s = DefaultSegment(r, 0x2e0)
         segments_copy[0:0] = [s]

atrcopy/dos33.py

@@ -598,6 +598,14 @@ class Dos33DiskImage(DiskImageBase):
             last = max(last, s.start_addr + len(s))
             log.debug("contiguous bytes needed: %04x - %04x" % (origin, last))
         if run_addr and run_addr != origin:
+            # check if run_addr points to some location that has data
+            found = False
+            for s in segments:
+                if run_addr >= s.start_addr and run_addr < s.start_addr + len(s):
+                    found = True
+                    break
+            if not found:
+                raise InvalidBinaryFile("Run address points outside data segments")
             origin -= 3
             hi, lo = divmod(run_addr, 256)
             raw = SegmentData([0x4c, lo, hi])
@@ -612,7 +620,7 @@ class Dos33DiskImage(DiskImageBase):
         words[1] = size
         for s in all_segments:
             index = s.start_addr - origin + 4
-            print "setting data for %04x - %04x at index %04x" % (s.start_addr, s.start_addr + len(s), index)
+            print "setting data for $%04x - $%04x at index $%04x" % (s.start_addr, s.start_addr + len(s), index)
             image[index:index + len(s)] = s.data
         return image, 'B'
 

test/test_create.py

@@ -0,0 +1,82 @@
+import numpy as np
+
+from mock import *
+
+from atrcopy import SegmentData, AtariDosDiskImage, Dos33DiskImage,InvalidBinaryFile, DefaultSegment
+from atrcopy.errors import *
+
+
+def get_image(file_name, diskimage_type):
+    data = np.fromfile(file_name, dtype=np.uint8)
+    rawdata = SegmentData(data)
+    image = diskimage_type(rawdata)
+    return image
+
+
+class BaseCreateTest(object):
+    diskimage_type = None
+
+    def get_exe_segments(self):
+        data1 = np.arange(4096, dtype=np.uint8)
+        data1[1::2] = np.repeat(np.arange(16, dtype=np.uint8), 128)
+        data2 = np.arange(4096, dtype=np.uint8)
+        data2[0::4] = np.repeat(np.arange(8, dtype=np.uint8), 128)
+        raw = [
+            (data1, 0x4000),
+            (data2, 0x8000),
+        ]
+
+        segments = []
+        for data, origin in raw:
+            rawdata = SegmentData(data)
+            s = DefaultSegment(rawdata, origin)
+            segments.append(s)
+        return segments
+
+    def check_exe(self, sample_file, diskimage_type, run_addr, expected):
+        image = get_image(sample_file, diskimage_type)
+        segments = self.get_exe_segments()
+        try:
+            _ = issubclass(AtrError, expected)
+            with pytest.raises(InvalidBinaryFile) as e:
+                file_data, filetype = image.create_executable_file_image(segments, run_addr)
+        except TypeError:
+            file_data, filetype = image.create_executable_file_image(segments, run_addr)
+            print image
+            print file_data, filetype
+            assert len(file_data) == expected
+
+@pytest.mark.parametrize("sample_file", ["../test_data/dos_sd_test1.atr"])
+class TestAtariDosSDImage(BaseCreateTest):
+    diskimage_type = AtariDosDiskImage
+
+    @pytest.mark.parametrize("run_addr,expected", [
+        (0x2000, InvalidBinaryFile),
+        (None, (2 + 6 + (4 + 0x1000) + (4 + 0x1000))),
+        (0x4000, (2 + 6 + (4 + 0x1000) + (4 + 0x1000))),
+        (0x8000, (2 + 6 + (4 + 0x1000) + (4 + 0x1000))),
+        (0xffff, InvalidBinaryFile),
+        ])
+    def test_exe(self, run_addr, expected, sample_file):
+        self.check_exe(sample_file, self.diskimage_type, run_addr, expected)
+
+
+@pytest.mark.parametrize("sample_file", ["../test_data/dos33_master.dsk"])
+class TestDos33Image(BaseCreateTest):
+    diskimage_type = Dos33DiskImage
+
+    @pytest.mark.parametrize("run_addr,expected", [
+        (0x2000, InvalidBinaryFile),
+        (None, (4 + (0x9000 - 0x4000))),
+        (0x4000, (4 + (0x9000 - 0x4000))),
+        (0x8000, (4 + 3 + (0x9000 - 0x4000))),
+        (0xffff, InvalidBinaryFile),
+        ])
+    def test_exe(self, run_addr, expected, sample_file):
+        self.check_exe(sample_file, self.diskimage_type, run_addr, expected)
+
+
+if __name__ == "__main__":
+    t = TestAtariDosSDImage()
+    t.setup()
+    t.test_exe()