From 9a3363a17f1f05cdfa71dd8914e5e4ed07f73c47 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89loi=20Rivard?= <eloi.rivard@nubla.fr>
Date: Sun, 11 Dec 2022 22:49:32 +0100
Subject: [PATCH] unit tests: improved flaskutils coverage

---
 canaille/flaskutils.py     | 11 +++++++++++
 canaille/oidc/endpoints.py | 16 +---------------
 tests/test_flaskutils.py   | 18 ++++++++++++++++++
 3 files changed, 30 insertions(+), 15 deletions(-)
 create mode 100644 tests/test_flaskutils.py

diff --git a/canaille/flaskutils.py b/canaille/flaskutils.py
index c8236885..78626859 100644
--- a/canaille/flaskutils.py
+++ b/canaille/flaskutils.py
@@ -1,6 +1,8 @@
 import datetime
 import logging
 from functools import wraps
+from urllib.parse import urlsplit
+from urllib.parse import urlunsplit
 
 import ldap
 from canaille.models import User
@@ -79,3 +81,12 @@ def smtp_needed():
 
 def timestamp(dt):
     return datetime.datetime.timestamp(dt)
+
+
+def set_parameter_in_url_query(url, **kwargs):
+    split = list(urlsplit(url))
+    pairs = split[3].split("&")
+    parameters = {pair.split("=")[0]: pair.split("=")[1] for pair in pairs if pair}
+    parameters = {**parameters, **kwargs}
+    split[3] = "&".join(f"{key}={value}" for key, value in parameters.items())
+    return urlunsplit(split)
diff --git a/canaille/oidc/endpoints.py b/canaille/oidc/endpoints.py
index e2b62f73..cad60455 100644
--- a/canaille/oidc/endpoints.py
+++ b/canaille/oidc/endpoints.py
@@ -1,6 +1,4 @@
 import datetime
-from urllib.parse import urlsplit
-from urllib.parse import urlunsplit
 
 from authlib.integrations.flask_oauth2 import current_token
 from authlib.jose import JsonWebKey
@@ -21,6 +19,7 @@ from flask_themer import render_template
 from werkzeug.datastructures import CombinedMultiDict
 
 from ..flaskutils import current_user
+from ..flaskutils import set_parameter_in_url_query
 from ..forms import FullLoginForm
 from ..models import User
 from .forms import LogoutForm
@@ -255,19 +254,6 @@ def userinfo():
     return jsonify(response)
 
 
-def set_parameter_in_url_query(url, **kwargs):
-    split = list(urlsplit(url))
-
-    parameters = "&".join(f"{key}={value}" for key, value in kwargs.items())
-
-    if split[3]:
-        split[3] = f"{split[3]}&{parameters}"
-    else:
-        split[3] = parameters
-
-    return urlunsplit(split)
-
-
 @bp.route("/end_session", methods=["GET", "POST"])
 def end_session():
     data = CombinedMultiDict((request.args, request.form))
diff --git a/tests/test_flaskutils.py b/tests/test_flaskutils.py
new file mode 100644
index 00000000..9bd9eddb
--- /dev/null
+++ b/tests/test_flaskutils.py
@@ -0,0 +1,18 @@
+from canaille.flaskutils import set_parameter_in_url_query
+
+
+def test_set_parameter_in_url_query():
+    assert (
+        set_parameter_in_url_query("https://auth.mydomain.tld", foo="bar")
+        == "https://auth.mydomain.tld?foo=bar"
+    )
+
+    assert (
+        set_parameter_in_url_query("https://auth.mydomain.tld?foo=baz", foo="bar")
+        == "https://auth.mydomain.tld?foo=bar"
+    )
+
+    assert (
+        set_parameter_in_url_query("https://auth.mydomain.tld?foo=baz", hello="world")
+        == "https://auth.mydomain.tld?foo=baz&hello=world"
+    )