update tests

.gitignore

@@ -1,3 +1,6 @@
+# idea stuff
+.idea/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

src/markdown_headers.py

@@ -1,12 +1,23 @@
+from collections import OrderedDict
 
 class MarkdownHeader:
+    head = "="
+    config = {
+        '======': 1,
+        '=====': 2,
+        '====': 3,
+        '===': 4,
+        '==': 5
+    }
 
     def convert(self, text):
-        if not text.startswith("="):
-            return text
-        elif text.startswith("== "):
-            return "######" + self.strip(text)
-        return "##"
+        config = OrderedDict(sorted(MarkdownHeader.config.items(), key = lambda t : t[1]))
+
+        for key, val in config.items():
+            if text.startswith(key):
+                return ('#' * val) + self.strip(text)
+        return text
+
 
     def strip(self, text):
-        return text.replace("=", "")
+        return text.replace(MarkdownHeader.head, "")

test/test_markdown_header.py

@@ -13,24 +13,20 @@ class TestMarkdownHeader(TestCase):
 
     def test_convert_h1(self):
         result = self.converter.convert("====== Classes ======")
-        self.assertEqual(result, "# Classes")
+        self.assertEqual(result, "# Classes ")
 
     def test_convert_h2(self):
         result = self.converter.convert("===== Classes =====")
-        self.assertEqual(result, "## Classes")
+        self.assertEqual(result, "## Classes ")
 
     def test_convert_h3(self):
         result = self.converter.convert("==== Classes ====")
-        self.assertEqual(result, "### Classes")
+        self.assertEqual(result, "### Classes ")
 
     def test_convert_h4(self):
-        result = self.converter.convert("==== Classes ====")
-        self.assertEqual(result, "#### Classes")
+        result = self.converter.convert("=== Classes ===")
+        self.assertEqual(result, "#### Classes ")
 
     def test_convert_h5(self):
-        result = self.converter.convert("=== Classes ===")
-        self.assertEqual(result, "##### Classes")
-
-    def test_convert_h6(self):
         result = self.converter.convert("== Classes ==")
-        self.assertEqual(result, "###### Classes ")
+        self.assertEqual(result, "##### Classes ")