From 20feb9c96a15d2e6e10cfc651d061d1d868ce90a Mon Sep 17 00:00:00 2001
From: Rob McMullen <feedback@playermissile.com>
Date: Fri, 3 Jun 2016 13:05:16 -0700
Subject: [PATCH] Added checks for invalid cart type and added unit tests

---
 atrcopy/__init__.py  |  2 +-
 atrcopy/cartridge.py |  9 +++--
 test/test_cart.py    | 81 ++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 4 deletions(-)
 create mode 100644 test/test_cart.py

diff --git a/atrcopy/__init__.py b/atrcopy/__init__.py
index 29344ad..63f08b8 100644
--- a/atrcopy/__init__.py
+++ b/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
 
diff --git a/atrcopy/cartridge.py b/atrcopy/cartridge.py
index d6500f2..4c44b51 100644
--- a/atrcopy/cartridge.py
+++ b/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]:
diff --git a/test/test_cart.py b/test/test_cart.py
new file mode 100644
index 0000000..c16383d
--- /dev/null
+++ b/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()
+