From 04e4da24b4c0cf78ccfd98e9950615c4482c2ac8 Mon Sep 17 00:00:00 2001
From: wgroeneveld <wouter.groeneveld@gmail.com>
Date: Thu, 12 Jan 2017 18:02:09 +0100
Subject: [PATCH] update tests

---
 .gitignore                   |  3 +++
 src/markdown_headers.py      | 23 +++++++++++++++++------
 test/test_markdown_header.py | 16 ++++++----------
 3 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/.gitignore b/.gitignore
index 72364f9..05bfe41 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,6 @@
+# idea stuff
+.idea/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
diff --git a/src/markdown_headers.py b/src/markdown_headers.py
index 44ed197..1c6b9ff 100644
--- a/src/markdown_headers.py
+++ b/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, "")
diff --git a/test/test_markdown_header.py b/test/test_markdown_header.py
index 356061a..8113553 100644
--- a/test/test_markdown_header.py
+++ b/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 ")