Refactoring and add testing

.cspell.json

@@ -3,6 +3,7 @@
     "version": "0.2",
     "words": [
         "actdiag",
+        "addopts",
         "Ashish",
         "autolayout",
         "backrefs",
@@ -31,6 +32,7 @@
         "Filebeat",
         "fontawesome",
         "fonttools",
+        "gantt",
         "graphviz",
         "Hideyuki",
         "hkato",
@@ -61,20 +63,24 @@
         "pypa",
         "pyphen",
         "pypi",
+        "pytest",
         "pyyaml",
         "rackdiag",
+        "ranksep",
         "remotedb",
         "Roboto",
         "sdist",
         "seqdiag",
         "Sigstore",
+        "skinparam",
         "soupsieve",
         "startuml",
         "structurizr",
         "superfences",
         "svgbob",
         "symbolator",
         "terrastruct",
+        "testpaths",
         "tikz",
         "tinycss",
         "urllib",

.github/workflows/publish-to-pypi.yml

@@ -14,6 +14,9 @@ jobs:
       - name: Install uv
         uses: astral-sh/setup-uv@v5
 
+      - name: Unit test
+        run: uv run pytest
+
       - name: Build the project
         run: uv build
 

pyproject.toml

@@ -42,6 +42,8 @@ dev = [
     "mkdocs-material>=9.6.11",
     "mkdocs-to-pdf>=0.10.0",
     "pymdown-extensions>=10.14.3",
+    "pytest>=8.3.5",
+    "pytest-mock>=3.14.0",
     "types-markdown>=3.8.0.20250415",
     "types-requests>=2.32.0.20250328",
 ]
@@ -50,3 +52,7 @@ dev = [
 line-length = 119
 [tool.ruff.format]
 quote-style = "single"
+
+[tool.pytest.ini_options]
+addopts = "-vv"
+testpaths = ["tests"]

src/markdown_kroki/extension.py

@@ -1,14 +1,13 @@
 """Diagram extension for Python-Markdown using Kroki"""
 
-import base64
 import re
 from typing import Generator, List
-import zlib
 
-import requests
 from markdown import Extension
 from markdown.preprocessors import Preprocessor
 
+from .kroki import KrokiServer
+
 
 class KrokiDiagramProcessor(Preprocessor):
     """Preprocessor to convert diagram code blocks to SVG/PNG image Data URIs."""
@@ -49,10 +48,6 @@ class KrokiDiagramProcessor(Preprocessor):
 
     KROKI_URL = 'https://kroki.io'
 
-    MIME_TYPES = {
-        'svg': 'image/svg+xml',
-        'png': 'image/png',
-    }
     IMG_TAG_ATTRIBUTES = [
         'alt',
         'width',
@@ -65,8 +60,9 @@ class KrokiDiagramProcessor(Preprocessor):
 
     def __init__(self, md, config):
         super().__init__(md)
-        self.kroki_url = config.get('kroki_url', self.KROKI_URL)
-        self.img_src = config.get('img_src', 'data')
+        kroki_url = config.get('kroki_url', self.KROKI_URL)
+        img_src = config.get('img_src', 'data')
+        self.kroki_server = KrokiServer(kroki_url, img_src)
 
     def run(self, lines: List[str]) -> List[str]:
         return list(self._parse_diagram_block(lines))
@@ -125,12 +121,13 @@ def _diagram_block_to_html(self, lines: List[str]) -> str:
                     if option in self.IMG_TAG_ATTRIBUTES:
                         img_tag_attributes[option] = code_block_options[option]
                     else:
-                        key = self._convert_mermaid_options_key(option)
                         kroki_options[key] = code_block_options[option].strip('"')
 
-                img_src = self._get_img_src(diagram_code, language, format, kroki_options)
+                # Get img src from Kroki server
+                img_src = self.kroki_server.get_img_src(diagram_code, language, format, kroki_options)
+
+                # Build the <img> tag with extracted attributes
                 if img_src:
-                    # Build the <img> tag with extracted options
                     img_tag = f'<img src="{img_src}"'
                     for key, value in img_tag_attributes.items():
                         img_tag += f' {key}={value}'
@@ -143,62 +140,6 @@ def _diagram_block_to_html(self, lines: List[str]) -> str:
 
         return html_string
 
-    def _convert_mermaid_options_key(self, key: str) -> str:
-        """Convert Mermaid options key to Kroki options key."""
-        # https://docs.kroki.io/kroki/setup/diagram-options/#_mermaid
-        key = re.sub('(.)([A-Z][a-z]+)', r'\1-\2', key)
-        key = re.sub('([a-z0-9])([A-Z])', r'\1-\2', key)
-        key = key.replace('.', '_').lower()
-        return key
-
-    def _get_img_src(self, diagram_code: str, language: str, format: str, kroki_options: dict) -> str:
-        """Convert diagram code to <img src>"""
-        if self.img_src == 'data':
-            # data URI
-            img_src = self._get_img_src_data(diagram_code, language, format, kroki_options)
-        elif self.img_src == 'link':
-            # Direct link
-            img_src = self._get_img_src_link(diagram_code, language, format, kroki_options)
-        return img_src
-
-    def _get_img_src_data(self, diagram_code: str, language: str, format: str, kroki_options: dict) -> str:
-        """Convert diagram code to img src data URI"""
-        base64image = self._get_base64image(diagram_code, language, format, kroki_options)
-        img_src = f'data:{self.MIME_TYPES[format]};base64,{base64image}'
-        return img_src
-
-    def _get_img_src_link(self, diagram_code: str, language: str, format: str, kroki_options: dict) -> str:
-        """Convert diagram code to img src direct link"""
-        encoded_code = base64.urlsafe_b64encode(zlib.compress(diagram_code.encode('utf-8'), 9)).decode('ascii')
-        option_list = []
-        options = ''
-        for key, value in kroki_options.items():
-            option_list.append(f'{key}={value}')
-        if option_list:
-            options = '?' + '&'.join(option_list)
-        # Kroki GET API
-        kroki_url = f'{self.kroki_url}/{language}/{format}/{encoded_code}{options}'
-        return kroki_url
-
-    def _get_base64image(self, diagram_code: str, language: str, format: str, kroki_options: dict) -> str:
-        """Convert diagram code to SVG/PNG using Kroki."""
-        kroki_url = f'{self.kroki_url}/{language}/{format}'
-        headers = {'Content-Type': 'text/plain'}
-        for key, value in kroki_options.items():
-            headers[f'Kroki-Diagram-Options-{key}'] = value
-
-        response = requests.post(kroki_url, headers=headers, data=diagram_code, timeout=30)
-        if response.status_code == 200:
-            if format == 'svg':
-                body = response.content.decode('utf-8')
-                base64image = base64.b64encode(body.encode('utf-8')).decode('utf-8')
-                return base64image
-            if format == 'png':
-                body = response.content
-                base64image = base64.b64encode(body).decode('utf-8')
-                return base64image
-        return ''
-
 
 class KrokiDiagramExtension(Extension):
     """Markdown Extension to support diagrams using Kroki."""

src/markdown_kroki/kroki.py

@@ -0,0 +1,99 @@
+"""
+KrokiServer is a class that provides methods to generate diagrams from text using the Kroki API.
+It supports various diagram types and formats, including SVG and PNG.
+It can generate data URIs or direct links to the generated diagrams.
+"""
+
+import base64
+import re
+import zlib
+
+import requests
+
+
+class KrokiServer:
+    """
+    KrokiServer is a class that represents a Kroki server.
+    It provides methods to generate diagrams from text using the Kroki API.
+    """
+
+    MIME_TYPES = {
+        'svg': 'image/svg+xml',
+        'png': 'image/png',
+    }
+
+    def __init__(self, kroki_url: str, img_src: str):
+        """ "
+        Initialize the KrokiServer with the given URL and image source type.
+        :param kroki_url: The URL of the Kroki server.
+        :param img_src: The type of image source ('data' or 'link').
+        """
+        self.kroki_url = kroki_url
+        self.img_src = img_src
+
+    def get_img_src(self, diagram_code: str, language: str, format: str, kroki_options: dict) -> str:
+        """
+        Generate the image source URL or data URI for the given diagram code.
+        :param diagram_code: The code for the diagram.
+        :param language: The language of the diagram (e.g., 'mermaid', 'plantuml').
+        :param format: The format of the image (e.g., 'svg', 'png').
+        :param kroki_options: Additional options for the Kroki API.
+        :return: The image source URL or data URI.
+        """
+        kroki_options = self._convert_mermaid_options_key(kroki_options)
+        if self.img_src == 'data':
+            # data URI
+            img_src = self._get_img_src_data(diagram_code, language, format, kroki_options)
+        elif self.img_src == 'link':
+            # Direct link
+            img_src = self._get_img_src_link(diagram_code, language, format, kroki_options)
+        return img_src
+
+    def _get_img_src_data(self, diagram_code: str, language: str, format: str, kroki_options: dict) -> str:
+        """Convert diagram code to img src data URI"""
+        base64image = self._get_base64image(diagram_code, language, format, kroki_options)
+        img_src = f'data:{self.MIME_TYPES[format]};base64,{base64image}'
+        return img_src
+
+    def _get_img_src_link(self, diagram_code: str, language: str, format: str, kroki_options: dict) -> str:
+        """Convert diagram code to img src link"""
+        encoded_code = base64.urlsafe_b64encode(zlib.compress(diagram_code.encode('utf-8'), 9)).decode('ascii')
+        option_list = []
+        options = ''
+        for key, value in kroki_options.items():
+            option_list.append(f'{key}={value}')
+        if option_list:
+            options = '?' + '&'.join(option_list)
+        # Kroki GET API
+        kroki_url = f'{self.kroki_url}/{language}/{format}/{encoded_code}{options}'
+        return kroki_url
+
+    def _get_base64image(self, diagram_code: str, language: str, format: str, kroki_options: dict) -> str:
+        """Convert diagram code to base64 image."""
+        kroki_url = f'{self.kroki_url}/{language}/{format}'
+        headers = {'Content-Type': 'text/plain'}
+        for key, value in kroki_options.items():
+            headers[f'Kroki-Diagram-Options-{key}'] = value
+
+        response = requests.post(kroki_url, headers=headers, data=diagram_code, timeout=30)
+        if response.status_code == 200:
+            if format == 'svg':
+                body = response.content.decode('utf-8')
+                base64image = base64.b64encode(body.encode('utf-8')).decode('utf-8')
+                return base64image
+            if format == 'png':
+                body = response.content
+                base64image = base64.b64encode(body).decode('utf-8')
+                return base64image
+        return ''
+
+    def _convert_mermaid_options_key(self, options: dict) -> dict:
+        """Convert Mermaid options key to Kroki options key."""
+        # https://docs.kroki.io/kroki/setup/diagram-options/#_mermaid
+        converted_options = {}
+        for key, value in options.items():
+            key = re.sub('(.)([A-Z][a-z]+)', r'\1-\2', key)
+            key = re.sub('([a-z0-9])([A-Z])', r'\1-\2', key)
+            key = key.replace('.', '_').lower()
+            converted_options[key] = value
+        return converted_options

tests/test_kroki.py

@@ -0,0 +1,110 @@
+"""Test KrokiServer class."""
+
+import pytest
+import requests
+
+from markdown_kroki.kroki import KrokiServer
+
+test_diagram_plantuml = """skinparam ranksep 20
+skinparam dpi 125
+skinparam packageTitleAlignment left
+
+rectangle "Main" {
+  (main.view)
+  (singleton)
+}
+rectangle "Base" {
+  (base.component)
+  (component)
+  (model)
+}
+rectangle "<b>main.ts</b>" as main_ts
+
+(component) ..> (base.component)
+main_ts ==> (main.view)
+(main.view) --> (component)
+(main.view) ...> (singleton)
+(singleton) ---> (model)"""
+
+test_diagram_mermaid = """graph TD
+  A[ Anyone ] -->|Can help | B( Go to github.com/yuzutech/kroki )
+  B --> C{ How to contribute? }
+  C --> D[ Reporting bugs ]
+  C --> E[ Sharing ideas ]
+  C --> F[ Advocating ]
+"""
+
+
+@pytest.mark.parametrize(
+    'diagram_code, language, format, options, expected_url',
+    [
+        (
+            test_diagram_plantuml,
+            'plantuml',
+            'svg',
+            {'theme': 'forest'},
+            'data:image/svg+xml;base64,ZmFrZV9pbWFnZV9kYXRh',
+        ),
+        (
+            test_diagram_mermaid,
+            'mermaid',
+            'png',
+            {'theme': 'base'},
+            'data:image/png;base64,ZmFrZV9pbWFnZV9kYXRh',
+        ),
+    ],
+)
+def test_get_img_src_data(mocker, diagram_code, language, format, options, expected_url):
+    """Test the get_img_src method."""
+    # Mock the requests.post method to return a fake response
+    mock_response = mocker.Mock(spec=requests.Response)
+    mock_response.content = b'fake_image_data'  # echo -n 'fake_image_data' | base64
+    mock_response.status_code = 200
+    mocker.patch('requests.post', return_value=mock_response)
+
+    kroki = KrokiServer('https://kroki.io', 'data')
+    result = kroki.get_img_src(diagram_code, language, format, options)
+    assert result == expected_url, f'Expected data URI, but got {result}'
+
+
+@pytest.mark.parametrize(
+    'diagram_code, language, format, options, expected_url',
+    [
+        (
+            test_diagram_plantuml,
+            'plantuml',
+            'svg',
+            {'theme': 'sketchy-outline'},
+            'https://kroki.io/plantuml/svg/eNpljzEPgjAQhff-iguTDFQlcYMmuru5mwNO0tCWhjY6GP-7LRJTdHvv7r67d26QxuKEGiY0gyML5Y65b7GzEvblIalYbAfs6SK9oqOSvdFkPCi6ecYmaj2aXhFkZ5QmgycD2Ogg-V3SI4_OyTjgR5OzVwqc0NECNEHydtR2NGH3TK2dHjtSP3zViPmQd9W2ERmgg-iv3jGW4MC5-L-wTEJdi1XeRENRiFWOtMfnrclriQ5gJD-Z3x9beAM=?theme=sketchy-outline',
+        ),
+        (
+            test_diagram_mermaid,
+            'mermaid',
+            'png',
+            {'theme': 'forest'},
+            'https://kroki.io/mermaid/png/eNpNzr0OgjAcBPCdp7hRB-QNNHz4Masb6VBK0zZg_6S2GhTf3cJgnH93l1OODxrXKgHyGrkdyUowpOl2KrmFlv2ACcUKR4InKON1aDaCbtkYXsFLobPOUWewjgvF3EP5xomec1qQ9c40MbbDJ3q5eFXjLAdy3liFJqg72M_2NS6au1lMK_k_HeK99kGCLz2WfAHnUzg5?theme=forest',
+        ),
+    ],
+)
+def test_get_img_src_link(diagram_code, language, format, options, expected_url):
+    """Test the get_img_src method."""
+    kroki = KrokiServer('https://kroki.io', 'link')
+    result = kroki.get_img_src(diagram_code, language, format, options)
+    assert result == expected_url, f'Expected data URI, but got {result}'
+
+
+@pytest.mark.parametrize(
+    'input_options, expected_options',
+    [
+        ({'theme': 'base'}, {'theme': 'base'}),
+        ({'fontFamily': 'courier'}, {'font-family': 'courier'}),
+        ({'maxTextSize': '50'}, {'max-text-size': '50'}),
+        ({'er.titleTopMargin': '100'}, {'er_title-top-margin': '100'}),
+        ({'gantt.displayMode': 'compact'}, {'gantt_display-mode': 'compact'}),
+    ],
+)
+def test_convert_mermaid_options_key(input_options, expected_options):
+    """Test the conversion of mermaid options keys."""
+    kroki = KrokiServer('https://kroki.io', 'data')
+    result = kroki._convert_mermaid_options_key(input_options)
+    assert result == expected_options, f'Expected {expected_options}, but got {result}'