dev(hansbug): add planutml unittest

Author: HansBug

.github/workflows/doc.yml

@@ -47,7 +47,7 @@ jobs:
           git branch -av
           git remote -v
           git tag
-          plantumlcli -c
+          python -m plantumlcli -c
           make pdocs
           mv ./docs/build/html ./public
       - name: Deploy to Github Page

plantumlcli/__main__.py

@@ -0,0 +1,4 @@
+from plantumlcli.entry.cli import cli as plantumlcli_cli
+
+if __name__ == '__main__':
+    plantumlcli_cli()

plantumlcli/models/remote.py

@@ -5,11 +5,11 @@
 import zlib
 from typing import Optional, Mapping, Any, Union, Tuple
 
-import requests
 from pyquery import PyQuery
 from urlobject import URLObject
 
 from .base import Plantuml, PlantumlResourceType, _has_cairosvg
+from ..utils import get_requests_session
 
 PLANTUML_HOST_ENV = 'PLANTUML_HOST'
 OFFICIAL_PLANTUML_HOST = 'http://www.plantuml.com/plantuml'
@@ -56,7 +56,7 @@ def __init__(self, host: str, **kwargs):
         _check_remote(self.__host)
         self.__host = _host_process(self.__host)
 
-        self.__session = requests.session()
+        self.__session = get_requests_session()
         self.__request_params = kwargs
 
     @classmethod

test/utils/test_session.py

@@ -0,0 +1,82 @@
+from unittest.mock import patch, Mock
+
+import pytest
+import requests
+from huggingface_hub import hf_hub_url
+from requests.adapters import HTTPAdapter
+
+from plantumlcli.utils.session import TimeoutHTTPAdapter, get_requests_session, get_random_ua
+
+
+@pytest.fixture
+def mock_requests_session():
+    with patch('requests.session') as mock_session:
+        yield mock_session.return_value
+
+
+@pytest.fixture
+def mock_ua_pool():
+    with patch('plantumlcli.utils.session._ua_pool') as mock_pool:
+        mock_pool.return_value.get_random_user_agent.return_value = 'MockUserAgent'
+        yield mock_pool
+
+
+@pytest.fixture()
+def example_url():
+    return hf_hub_url(
+        repo_id='deepghs/danbooru_newest',
+        repo_type='dataset',
+        filename='README.md'
+    )
+
+
+@pytest.mark.unittest
+class TestUtilsSession:
+    def test_timeout_http_adapter_init(self, ):
+        adapter = TimeoutHTTPAdapter()
+        assert adapter.timeout == 15
+
+        adapter = TimeoutHTTPAdapter(timeout=30)
+        assert adapter.timeout == 30
+
+    def test_timeout_http_adapter_send(self, ):
+        adapter = TimeoutHTTPAdapter(timeout=10)
+        mock_request = Mock()
+        mock_kwargs = {}
+
+        with patch.object(HTTPAdapter, 'send') as mock_send:
+            adapter.send(mock_request, **mock_kwargs)
+            mock_send.assert_called_once_with(mock_request, timeout=10)
+
+        mock_kwargs = {'timeout': 20}
+        with patch.object(HTTPAdapter, 'send') as mock_send:
+            adapter.send(mock_request, **mock_kwargs)
+            mock_send.assert_called_once_with(mock_request, timeout=20)
+
+    def test_get_requests_session(self, mock_ua_pool):
+        session = get_requests_session()
+        assert isinstance(session, requests.Session)
+        assert 'User-Agent' in session.headers
+        assert session.headers['User-Agent'] == 'MockUserAgent'
+
+        custom_headers = {'Custom-Header': 'Value'}
+        session = get_requests_session(headers=custom_headers)
+        assert 'Custom-Header' in session.headers
+        assert session.headers['Custom-Header'] == 'Value'
+
+        session = get_requests_session(verify=False)
+        assert session.verify is False
+
+        existing_session = requests.Session()
+        session = get_requests_session(session=existing_session)
+        assert session is existing_session
+
+    def test_get_requests_session_with_custom_params(self):
+        session = get_requests_session(max_retries=3, timeout=30)
+        assert isinstance(session, requests.Session)
+        # You might want to add more assertions here to check if the custom parameters are applied correctly
+
+    def test_get_random_ua(self, mock_ua_pool):
+        ua = get_random_ua()
+        assert ua == 'MockUserAgent'
+        mock_ua_pool.return_value.get_random_user_agent.assert_called_once()