Add support for additional request parameters in login methods (#865)

* Add support for additional request parameters in login methods

* Refactor tests to use pytest.raises instead of unittest assertions

* Update wikibaseintegrator/wbi_login.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Use filtered kwargs for all queries

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: LeMyst

test/test_wbi_helpers.py

@@ -2,6 +2,7 @@
 import os
 import unittest
 
+import pytest
 import requests
 
 from wikibaseintegrator.wbi_config import config as wbi_config
@@ -15,22 +16,22 @@ def test_connection():
 
     mediawiki_api_call_helper(data=data, max_retries=2, retry_after=1, allow_anonymous=True)
 
-    with unittest.TestCase().assertRaises(MaxRetriesReachedException):
+    with pytest.raises(MaxRetriesReachedException):
         mediawiki_api_call_helper(data=data, mediawiki_api_url="https://www.wikidataaaaaaa.org", max_retries=2, retry_after=1, allow_anonymous=True)
 
-    with unittest.TestCase().assertRaises(MaxRetriesReachedException):
+    with pytest.raises(MaxRetriesReachedException):
         mediawiki_api_call_helper(data=data, mediawiki_api_url=os.getenv("HTTPSTATUS_SERVICE", "https://httpbin.org") + "/status/500", max_retries=2, retry_after=1, allow_anonymous=True)
 
-    with unittest.TestCase().assertRaises(MaxRetriesReachedException):
+    with pytest.raises(MaxRetriesReachedException):
         mediawiki_api_call_helper(data=data, mediawiki_api_url=os.getenv("HTTPSTATUS_SERVICE", "https://httpbin.org") + "/status/502", max_retries=2, retry_after=1, allow_anonymous=True)
 
-    with unittest.TestCase().assertRaises(MaxRetriesReachedException):
+    with pytest.raises(MaxRetriesReachedException):
         mediawiki_api_call_helper(data=data, mediawiki_api_url=os.getenv("HTTPSTATUS_SERVICE", "https://httpbin.org") + "/status/503", max_retries=2, retry_after=1, allow_anonymous=True)
 
-    with unittest.TestCase().assertRaises(MaxRetriesReachedException):
+    with pytest.raises(MaxRetriesReachedException):
         mediawiki_api_call_helper(data=data, mediawiki_api_url=os.getenv("HTTPSTATUS_SERVICE", "https://httpbin.org") + "/status/504", max_retries=2, retry_after=1, allow_anonymous=True)
 
-    with unittest.TestCase().assertRaises(requests.HTTPError):
+    with pytest.raises(requests.HTTPError):
         mediawiki_api_call_helper(data=data, mediawiki_api_url=os.getenv("HTTPSTATUS_SERVICE", "https://httpbin.org") + "/status/400", max_retries=2, retry_after=1, allow_anonymous=True)
 
 
@@ -56,7 +57,7 @@ def test_user_agent(caplog):
 def test_allow_anonymous():
     wbi_config['USER_AGENT'] = 'WikibaseIntegrator-pytest/1.0 (test_wbi_helpers.py)'
     # Test there is a warning because of allow_anonymous
-    with unittest.TestCase().assertRaises(ValueError):
+    with pytest.raises(ValueError):
         mediawiki_api_call_helper(data={'format': 'json', 'action': 'wbgetentities', 'ids': 'Q42'}, max_retries=3, retry_after=1, user_agent='MyWikibaseBot/0.5')
 
     # Test there is no warning because of allow_anonymous
@@ -84,7 +85,7 @@ def test_format2wbi():
     assert isinstance(format2wbi('property', '{}'), PropertyEntity)
     assert isinstance(format2wbi('lexeme', '{}'), LexemeEntity)
     assert isinstance(format2wbi('mediainfo', '{}'), MediaInfoEntity)
-    with unittest.TestCase().assertRaises(ValueError):
+    with pytest.raises(ValueError):
         format2wbi('unknown', '{}')
 
     result = format2wbi('item', '''{

test/test_wbi_login.py

@@ -3,6 +3,7 @@
 import unittest
 
 import pytest
+import requests
 from oauthlib.oauth2 import MissingTokenError
 
 from wikibaseintegrator import wbi_login
@@ -23,11 +24,11 @@
 
 
 def test_login():
-    with unittest.TestCase().assertRaises(LoginError):
+    with pytest.raises(LoginError):
         login = wbi_login.Clientlogin(user='wrong', password='wrong')
         login.generate_edit_credentials()
 
-    with unittest.TestCase().assertRaises(LoginError):
+    with pytest.raises(LoginError):
         login = wbi_login.Login(user='wrong', password='wrong')
         login.generate_edit_credentials()
 
@@ -38,8 +39,16 @@ def test_login():
         print("no WDUSER or WDPASS found in environment variables", file=sys.stderr)
 
 
+def test_verify():
+    with pytest.raises(requests.exceptions.SSLError):
+        wbi_login.Clientlogin(user='wrong', password='wrong', mediawiki_api_url='https://self-signed.badssl.com/', verify=True)
+
+    with pytest.raises(requests.exceptions.JSONDecodeError):
+        wbi_login.Clientlogin(user='wrong', password='wrong', mediawiki_api_url='https://self-signed.badssl.com/', verify=False)
+
+
 def test_oauth1():
-    with unittest.TestCase().assertRaises(LoginError):
+    with pytest.raises(LoginError):
         login = wbi_login.OAuth1(consumer_token='wrong', consumer_secret='wrong')
         login.generate_edit_credentials()
 
@@ -50,7 +59,7 @@ def test_oauth1():
 
 
 def test_oauth1_access():
-    with unittest.TestCase().assertRaises(LoginError):
+    with pytest.raises(LoginError):
         login = wbi_login.OAuth1(consumer_token='wrong', consumer_secret='wrong', access_token='wrong', access_secret='wrong')
         login.generate_edit_credentials()
 

wikibaseintegrator/wbi_login.py

@@ -4,7 +4,7 @@
 import logging
 import time
 import webbrowser
-from typing import Optional
+from typing import Any, Optional
 
 from mwoauth import ConsumerToken, Handshaker, OAuthException
 from oauthlib.oauth2 import BackendApplicationClient, InvalidClientError
@@ -199,7 +199,7 @@ def continue_oauth(self, oauth_callback_data: Optional[str] = None) -> None:
 
 class Login(_Login):
     @wbi_backoff()
-    def __init__(self, user: Optional[str] = None, password: Optional[str] = None, mediawiki_api_url: Optional[str] = None, token_renew_period: int = 1800, user_agent: Optional[str] = None):
+    def __init__(self, user: Optional[str] = None, password: Optional[str] = None, mediawiki_api_url: Optional[str] = None, token_renew_period: int = 1800, user_agent: Optional[str] = None, **kwargs: Any):
         """
         This class is used to log in with a bot password
 
@@ -208,6 +208,7 @@ def __init__(self, user: Optional[str] = None, password: Optional[str] = None, m
         :param mediawiki_api_url: The URL to the MediaWiki API (default Wikidata)
         :param token_renew_period: Seconds after which a new token should be requested from the Wikidata server
         :param user_agent: UA string to use for API requests.
+        :param kwargs: Additional parameters to pass to the requests.sessions.Session.post method, such as headers or proxies.
         """
 
         mediawiki_api_url = str(mediawiki_api_url or config['MEDIAWIKI_API_URL'])
@@ -225,9 +226,13 @@ def __init__(self, user: Optional[str] = None, password: Optional[str] = None, m
             'User-Agent': get_user_agent(user_agent)
         }
 
-        # get login token
-        login_token = session.post(mediawiki_api_url, data=params_login, headers=headers).json()['query']['tokens']['logintoken']
+        allowed_kwargs = {'headers', 'proxies', 'timeout', 'verify'}
+        filtered_kwargs = {key: value for key, value in kwargs.items() if key in allowed_kwargs}
+        if len(filtered_kwargs) < len(kwargs):
+            log.warning("Unsupported kwargs were ignored: %s", set(kwargs) - allowed_kwargs)
 
+        # get login token
+        login_token = session.post(mediawiki_api_url, data=params_login, headers=headers, **filtered_kwargs).json()['query']['tokens']['logintoken']
         params = {
             'action': 'login',
             'lgname': user,
@@ -236,7 +241,7 @@ def __init__(self, user: Optional[str] = None, password: Optional[str] = None, m
             'format': 'json'
         }
 
-        login_result = session.post(mediawiki_api_url, data=params, headers=headers).json()
+        login_result = session.post(mediawiki_api_url, data=params, headers=headers, **filtered_kwargs).json()
 
         if 'login' in login_result and login_result['login']['result'] == 'Success':
             log.info("Successfully logged in as %s", login_result['login']['lgusername'])
@@ -253,7 +258,7 @@ def __init__(self, user: Optional[str] = None, password: Optional[str] = None, m
 
 class Clientlogin(_Login):
     @wbi_backoff()
-    def __init__(self, user: Optional[str] = None, password: Optional[str] = None, mediawiki_api_url: Optional[str] = None, token_renew_period: int = 1800, user_agent: Optional[str] = None):
+    def __init__(self, user: Optional[str] = None, password: Optional[str] = None, mediawiki_api_url: Optional[str] = None, token_renew_period: int = 1800, user_agent: Optional[str] = None, **kwargs: Any):
         """
         This class is used to log in with a user account
 
@@ -262,6 +267,7 @@ def __init__(self, user: Optional[str] = None, password: Optional[str] = None, m
         :param mediawiki_api_url: The URL to the MediaWiki API (default Wikidata)
         :param token_renew_period: Seconds after which a new token should be requested from the Wikidata server
         :param user_agent: UA string to use for API requests.
+        :param kwargs: Additional parameters to pass to the requests.sessions.Session.post method, such as headers or proxies.
         """
 
         mediawiki_api_url = str(mediawiki_api_url or config['MEDIAWIKI_API_URL'])
@@ -279,8 +285,13 @@ def __init__(self, user: Optional[str] = None, password: Optional[str] = None, m
             'User-Agent': get_user_agent(user_agent)
         }
 
+        allowed_kwargs = {'headers', 'proxies', 'timeout', 'verify'}
+        filtered_kwargs = {key: value for key, value in kwargs.items() if key in allowed_kwargs}
+        if len(filtered_kwargs) < len(kwargs):
+            log.warning("Unsupported kwargs were ignored: %s", set(kwargs) - allowed_kwargs)
+
         # get login token
-        login_token = session.post(mediawiki_api_url, data=params_login, headers=headers).json()['query']['tokens']['logintoken']
+        login_token = session.post(mediawiki_api_url, data=params_login, headers=headers, **filtered_kwargs).json()['query']['tokens']['logintoken']
 
         params = {
             'action': 'clientlogin',
@@ -291,7 +302,7 @@ def __init__(self, user: Optional[str] = None, password: Optional[str] = None, m
             'format': 'json'
         }
 
-        login_result = session.post(mediawiki_api_url, data=params, headers=headers).json()
+        login_result = session.post(mediawiki_api_url, data=params, headers=headers, **filtered_kwargs).json()
 
         log.debug(login_result)