From e1d7592e4307de2e1cb955037839ea1c61ee5def Mon Sep 17 00:00:00 2001
From: Ondrej Tuma <mcbig@zeropage.cz>
Date: Wed, 28 Aug 2024 14:06:22 +0200
Subject: [PATCH] Append MK(X)S printers for Connect

---
 prusaerrors/connect/codes.py |  7 +++--
 prusaerrors/shared/codes.py  |  5 +++-
 tests/test_connect.py        | 54 ++++++++++++++++++++++++++++++++++--
 3 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/prusaerrors/connect/codes.py b/prusaerrors/connect/codes.py
index afeaed5..d647581 100644
--- a/prusaerrors/connect/codes.py
+++ b/prusaerrors/connect/codes.py
@@ -16,7 +16,8 @@
 from prusaerrors.shared.codes import unique_codes, Codes, Printer, Code, Category
 
 
-BUDDY = ['MINI', 'MK4', 'IX', 'XL', 'MK35', 'MK39']
+BUDDY = ['MINI', 'MK4', 'IX', 'XL', 'MK3.5',
+         'MK4S', 'MK3.9', 'MK3.9S', 'MK3.5S']
 
 
 class PrinterCode(Code):
@@ -73,7 +74,9 @@ def decor(cls):
             if code_parts["printer"] == 'XX':
                 if printers := entry.get("printers"):
                     if 'MK4' in printers:
-                        printers.append('MK39')
+                        printers.extend(('MK4S', 'MK3.9', 'MK3.9S'))
+                    elif 'MK3.5' in printers:
+                        printers.append('MK3.5S')
                 else:  # if no printers specified code is valid for all buddy
                     printers = BUDDY
 
diff --git a/prusaerrors/shared/codes.py b/prusaerrors/shared/codes.py
index c7f6240..af162a7 100644
--- a/prusaerrors/shared/codes.py
+++ b/prusaerrors/shared/codes.py
@@ -30,8 +30,11 @@ class Printer(IntEnum):
     MK4 = 0x000D
     IX = 0x0010
     XL = 0x0011
-    MK35 = 0x0017
     MK39 = 0x0015
+    MK35 = 0x0017
+    MK4S = 0x001A
+    MK39S = 0x001B
+    MK35S = 0x001C
     M1 = 0x001D
 
 
diff --git a/tests/test_connect.py b/tests/test_connect.py
index 5cc8284..f304e60 100644
--- a/tests/test_connect.py
+++ b/tests/test_connect.py
@@ -14,15 +14,65 @@
 
 class TestErrors(unittest.TestCase):
 
-    def test_code_lookup(self):
+    def test_XL(self):
         code = PrinterCodes.get("17505")
-        assert code.printer == Printer(17)
+        assert code.printer == Printer.XL
         assert code.category == Category(5)
         assert code.error == 5
         assert code.title
         assert code.message
         assert code.id
 
+    def test_MK4(self):
+        code = PrinterCodes.get("13505")
+        assert code.printer == Printer.MK4
+        assert code.category == Category(5)
+        assert code.error == 5
+        assert code.title
+        assert code.message
+        assert code.id
+
+    def test_MK4S(self):
+        code = PrinterCodes.get("26505")
+        assert code.printer == Printer.MK4S
+        assert code.category == Category(5)
+        assert code.error == 5
+        assert code.title
+        assert code.message
+        assert code.id
+
+    def test_MK39(self):
+        code = PrinterCodes.get("21505")
+        assert code.printer == Printer.MK39
+        assert code.category == Category(5)
+        assert code.error == 5
+        assert code.title
+        assert code.message
+        assert code.id
+
+    def test_MK35(self):
+        code = PrinterCodes.get("23701")
+        assert code.printer == Printer.MK35
+        assert code.category == Category(7)
+        assert code.error == 1
+        assert code.title
+        assert code.message
+        assert code.id
+
+    def test_MK35S(self):
+        code = PrinterCodes.get("28701")
+        assert code.printer == Printer.MK35S
+        assert code.category == Category(7)
+        assert code.error == 1
+        assert code.title
+        assert code.message
+        assert code.id
+
+    def test_no_MK35S(self):
+        """MK35S doesn't have puppies."""
+        code = PrinterCodes.get("28512")
+        assert code is None
+
     def test_unknown_code(self):
         code = PrinterCodes.get("unknown_code")
         assert code is None