diff --git a/eorg/const.py b/eorg/const.py
index bdb1c1e..0d3e6f7 100755
--- a/eorg/const.py
+++ b/eorg/const.py
@@ -59,7 +59,7 @@ TOKENS = {
     ),
     tokens.IMAGE: TokenStruct(start=t_IMG, end_pos=-2),
     tokens.CAPTION: TokenStruct(
-        start=t_CAPTIONS, type=TYPE_ATTRIBUTE, key="CAPTION"
+        start=t_CAPTIONS, type=TYPE_ATTRIBUTE, key="caption"
     ),
     tokens.SOURCE: TokenStruct(start=t_SRC_BEGIN, end=t_SRC_END),
     tokens.TABLE: TokenStruct(
@@ -73,17 +73,6 @@ TOKENS = {
     tokens.META_OTHER: TokenStruct(
         start=t_META_OTHER, start_pos=2, end_pos=-1
     ),
-    # tokens.META: (t_META, False, 2, -1, False),
-    # tokens.COMMENT: (t_COMMENT_BEGIN, t_COMMENT_END, 2, None, False),
-    # tokens.EXAMPLE: (t_EXAMPLE_BEGIN, t_EXAMPLE_END, 2, None, False),
-    # tokens.IMAGE: (t_IMG, False, 2, None, False),
-    # tokens.CAPTION: (t_CAPTIONS, False, 2, None, False),
-    # tokens.SOURCE: (t_SRC_BEGIN, t_SRC_END, 2, None, False),
-    # tokens.TABLE: (t_TABLE_START, t_TABLE_END, 0, None, False),
-    # tokens.BULLET: (t_BULLET_START, t_BULLET_END, 0, None, False),
-    # tokens.RESULTS: (t_SRC_BEGIN, t_SRC_END, 2, None, False),
-    # tokens.HEADER: (t_HEADER, False, 1, None, True),
-    # tokens.META_OTHER: (t_META_OTHER, False, 2, -1, False),
 }
 
 
diff --git a/eorg/generate.py b/eorg/generate.py
index 5f52ac7..650d24a 100644
--- a/eorg/generate.py
+++ b/eorg/generate.py
@@ -5,6 +5,7 @@ from eorg.const import Token, ESCAPE
 from eorg import tokens
 from eorg.tokens import Token
 from pygments import highlight
+from pygments.util import ClassNotFound
 from pygments.lexers import PythonLexer
 from pygments.lexers import get_lexer_by_name
 from pygments.formatters import HtmlFormatter
@@ -13,7 +14,7 @@ from pygments.formatters import HtmlFormatter
 def src(doc, code, cls="", root=True):
     try:
         lexer = get_lexer_by_name(code.attrs.get("language", "shell"))
-    except pygments.util.ClassNotFound as e:
+    except ClassNotFound as e:
         lexer = get_lexer_by_name(code.attrs.get("language", "text"))
     return highlight(code.value, lexer, HtmlFormatter(linenos=True))
 
@@ -21,8 +22,10 @@ def src(doc, code, cls="", root=True):
 def img(doc, item, cls="", root=True):
     caption = doc.previous(tokens.CAPTION)
     text = ""
-    if caption:
-        text = f'<p class="center-align">{caption.value}</p>'
+    if item.attrs:
+        caption = item.attrs.get('caption')
+        if caption:
+            text = f'<p class="center-align">{caption}</p>'
     return f'<img{cls} style="margin:auto;" src="{item.value[0]}" alt="{item.value[1]}" />{text}'
 
 
diff --git a/setup.cfg b/setup.cfg
new file mode 100644
index 0000000..9af7e6f
--- /dev/null
+++ b/setup.cfg
@@ -0,0 +1,2 @@
+[aliases]
+test=pytest
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 89c6e62..650b3d8 100755
--- a/setup.py
+++ b/setup.py
@@ -49,5 +49,5 @@ setup(
 
     install_requires=["pygments"],
     setup_requires=['pytest-runner',],
-    tests_require=['pytest']
+    tests_require=['pytest', 'pytest-runner']
 )
diff --git a/tests/test_document_parsing.py b/tests/test_document_parsing.py
index 8f9d16b..6d9c6aa 100644
--- a/tests/test_document_parsing.py
+++ b/tests/test_document_parsing.py
@@ -67,7 +67,7 @@ def test_image_with_caption():
         Token(
             tokens.IMAGE,
             ["../../test.jpg", ""],
-            attrs={"CAPTION": " Test Image"},
+            attrs={"caption": " Test Image"},
         )
     ]
     result = parse(text).doc
@@ -84,7 +84,7 @@ def test_image_with_caption():
         Token(
             tokens.IMAGE,
             ["../../test.jpg", "test"],
-            attrs={"CAPTION": " Test Image"},
+            attrs={"caption": " Test Image"},
         )
     ]
     result = parse(text).doc