Added checks for invalid cart type and added unit tests

atrcopy/__init__.py

@@ -13,7 +13,7 @@ from diskimages import AtrHeader, BootDiskImage, add_atr_header
 from kboot import KBootImage, add_xexboot_header
 from segments import SegmentData, SegmentSaver, DefaultSegment, EmptySegment, ObjSegment, RawSectorsSegment, user_bit_mask, match_bit_mask, comment_bit_mask, data_bit_mask, selected_bit_mask, diff_bit_mask, not_user_bit_mask, interleave_segments
 from spartados import SpartaDosDiskImage
-from cartridge import A8CartHeader
+from cartridge import A8CartHeader, AtariCartImage
 from parsers import SegmentParser, DefaultSegmentParser, guess_parser_for_mime, guess_parser_for_system, iter_parsers, iter_known_segment_parsers, mime_parse_order
 from utils import to_numpy
 

atrcopy/cartridge.py

@@ -26,7 +26,7 @@ known_cart_types = [
     (0,  "", 0,),
     (57, "Standard 2 KB", 2, 2, 0, 0xb800),
     (58, "Standard 4 KB", 4, 4, 0, 0xb000),
-    (59, "Right slot 4 KB", 4, 4, 4, 0, 0x9000),
+    (59, "Right slot 4 KB", 4, 4, 0, 0, 0x9000),
     (1,  "Standard 8 KB", 8, 8, 0, 0xa000),
     (21, "Right slot 8 KB", 8,),
     (2,  "Standard 16 KB", 16, 16, 0, 0x8000),
@@ -103,7 +103,10 @@ def get_known_carts():
     return grouped
 
 def get_cart(cart_type):
-    return known_cart_types[known_cart_type_map[cart_type]]
+    try:
+        return known_cart_types[known_cart_type_map[cart_type]]
+    except KeyError:
+        raise InvalidDiskImage("Unsupported cart type %d" % cart_type)
 
 
 class A8CartHeader(object):
@@ -222,7 +225,7 @@ class AtariCartImage(DiskImageBase):
             return
         k, rem = divmod((len(self) - len(self.header)), 1024)
         c = get_cart(self.cart_type)
-        log.debug("checking %s:" % c[1], k, rem, c[2])
+        log.debug("checking type=%d, k=%d, rem=%d for %s, %s" % (self.cart_type, k, rem, c[1], c[2]))
         if rem > 0:
             raise InvalidDiskImage("Cart not multiple of 1K")
         if k != c[2]:

test/test_cart.py

@@ -0,0 +1,81 @@
+from mock import *
+
+from atrcopy import AtariCartImage, SegmentData, InvalidDiskImage
+
+
+class TestAtariCart(object):
+    def setup(self):
+        pass
+
+    def get_cart(self, k_size, cart_type):
+        data = np.zeros((k_size * 1024)+16, dtype=np.uint8)
+        data[0:4].view("|a4")[0] = 'CART'
+        data[4:8].view(">u4")[0] = cart_type
+        return data
+
+    def test_unbanked(self):
+        carts = [
+            (8, 1),
+            (16, 2),
+            (8, 21),
+            (2, 57),
+            (4, 58),
+            (4, 59),
+        ]
+        for k_size, cart_type in carts:
+            data = self.get_cart(k_size, cart_type)
+            rawdata = SegmentData(data)
+            image = AtariCartImage(rawdata, cart_type)
+            image.parse_segments()
+            assert len(image.segments) == 2
+            assert len(image.segments[0]) == 16
+            assert len(image.segments[1]) == k_size * 1024
+
+    def test_banked(self):
+        carts = [
+            (32, 8, 8, 12),
+            (64, 8, 8, 13),
+            (64, 8, 8, 67),
+            (128, 8, 8, 14),
+            (256, 8, 8, 23),
+            (512, 8, 8, 24),
+            (1024, 8, 8, 25),
+        ]
+        for k_size, main_size, banked_size, cart_type in carts:
+            data = self.get_cart(k_size, cart_type)
+            rawdata = SegmentData(data)
+            image = AtariCartImage(rawdata, cart_type)
+            image.parse_segments()
+            assert len(image.segments) == 1 + 1 + (k_size - main_size)/banked_size
+            assert len(image.segments[0]) == 16
+            assert len(image.segments[1]) == main_size * 1024
+            assert len(image.segments[2]) == banked_size * 1024
+
+    def test_bad(self):
+        k_size = 32
+
+        # check for error because invalid data in cart image itself
+        data = self.get_cart(k_size, 1337)
+        rawdata = SegmentData(data)
+        with pytest.raises(InvalidDiskImage):
+            image = AtariCartImage(rawdata, 1337)
+        with pytest.raises(InvalidDiskImage):
+            image = AtariCartImage(rawdata, 12)
+
+        # check for error with valid cart image, but invalid cart type supplied
+        # to the image parser
+        data = self.get_cart(k_size, 12)
+        rawdata = SegmentData(data)
+        with pytest.raises(InvalidDiskImage):
+            image = AtariCartImage(rawdata, 1337)
+
+
+
+
+if __name__ == "__main__":
+    print "\n".join(mime_parse_order)
+
+    t = TestAtariCart()
+    t.setup()
+    #t.test_segment()
+