feat: add tests for the AuthDB cleanup function

Author: HeloiseJoffe

diracx-db/tests/auth/test_authorization_flow.py

@@ -5,6 +5,7 @@
 
 from diracx.core.exceptions import AuthorizationError
 from diracx.db.sql.auth.db import AuthDB
+from diracx.db.sql.auth.schema import FlowStatus
 
 MAX_VALIDITY = 2
 EXPIRED = 0
@@ -74,3 +75,39 @@ async def test_insert(auth_db: AuthDB):
         )
 
     assert uuid1 != uuid2
+
+
+async def test_clean_authorization_flows(auth_db: AuthDB):
+    # Insert two authorization flows
+    async with auth_db as auth_db:
+        uuid1 = await auth_db.insert_authorization_flow(
+            "client_id", "scope", "code_challenge", "S256", "redirect_uri"
+        )
+        uuid2 = await auth_db.insert_authorization_flow(
+            "client_id2", "scope2", "code_challenge2", "S256", "redirect_uri2"
+        )
+
+    id_token = {"sub": "myIdToken"}
+
+    async with auth_db as auth_db:
+        code1, _ = await auth_db.authorization_flow_insert_id_token(uuid1, id_token, 1)
+        code2, _ = await auth_db.authorization_flow_insert_id_token(uuid2, id_token, 1)
+
+    async with auth_db as auth_db:
+        await auth_db.update_authorization_flow_status(code1, FlowStatus.DONE)
+        await auth_db.update_authorization_flow_status(code2, FlowStatus.ERROR)
+
+    # Check the number of deleted authorization flow (should be 0)
+    async with auth_db as auth_db:
+        deleted_auth = await auth_db.clean_expired_authorization_flows(max_retention=30)
+        assert deleted_auth == 0
+
+    # Check the number of deleted authorization flow (should be 2)
+    async with auth_db as auth_db:
+        deleted_auth = await auth_db.clean_expired_authorization_flows(max_retention=0)
+        assert deleted_auth == 2
+
+    # Check the number of deleted authorization flow (should be 0 because there is nothing left to delete)
+    async with auth_db as auth_db:
+        deleted_auth = await auth_db.clean_expired_authorization_flows(max_retention=0)
+        assert deleted_auth == 0

diracx-db/tests/auth/test_device_flow.py

@@ -8,7 +8,7 @@
 
 from diracx.core.exceptions import AuthorizationError
 from diracx.db.sql.auth.db import AuthDB
-from diracx.db.sql.auth.schema import USER_CODE_LENGTH
+from diracx.db.sql.auth.schema import USER_CODE_LENGTH, FlowStatus
 from diracx.db.sql.utils.functions import substract_date
 
 MAX_VALIDITY = 2
@@ -139,3 +139,33 @@ async def test_device_flow_insert_id_token(auth_db: AuthDB):
     async with auth_db as auth_db:
         res = await auth_db.get_device_flow(device_code)
         assert res["IDToken"] == id_token
+
+
+async def test_clean_device_flows(auth_db: AuthDB):
+    # Insert two device flows
+    async with auth_db as auth_db:
+        user_code1, device_code1 = await auth_db.insert_device_flow(
+            "client_id", "scope"
+        )
+        user_code2, device_code2 = await auth_db.insert_device_flow(
+            "client_id", "scope"
+        )
+
+    async with auth_db as auth_db:
+        await auth_db.update_device_flow_status(device_code1, FlowStatus.DONE)
+        await auth_db.update_device_flow_status(device_code2, FlowStatus.ERROR)
+
+    # Check the number of deleted device flows (should be 0)
+    async with auth_db as auth_db:
+        deleted_device = await auth_db.clean_expired_device_flows(max_retention=30)
+        assert deleted_device == 0
+
+    # Check the number of deleted device flows (should be 2)
+    async with auth_db as auth_db:
+        deleted_device = await auth_db.clean_expired_device_flows(max_retention=0)
+        assert deleted_device == 2
+
+    # Check the number of deleted device flow (should be 0 because there is nothing left to delete)
+    async with auth_db as auth_db:
+        deleted_device = await auth_db.clean_expired_device_flows(max_retention=0)
+        assert deleted_device == 0

diracx-db/tests/auth/test_refresh_token.py

@@ -257,3 +257,44 @@ async def test_get_refresh_tokens(auth_db: AuthDB):
 
     # Check the number of retrieved refresh tokens (should be 3 refresh tokens)
     assert len(refresh_tokens) == 2
+
+
+async def test_clean_refresh_tokens(auth_db: AuthDB):
+    # Insert two refresh tokens
+    jtis = []
+    async with auth_db as auth_db:
+        for _ in range(2):
+            jti = uuid7()
+            await auth_db.insert_refresh_token(
+                jti,
+                "subject",
+                "scope",
+            )
+            jtis.append(jti)
+
+    # Revoke one of the refresh token
+    async with auth_db as auth_db:
+        await auth_db.revoke_refresh_token(jtis[0])
+
+    # Check the number of deleted refresh tokens (should be 0)
+    async with auth_db as auth_db:
+        deleted_expired, deleted_revoked = await auth_db.clean_expired_refresh_token(
+            max_validity=10, max_retention=30
+        )
+    assert deleted_expired == 0
+    assert deleted_revoked == 0
+
+    # Check the number of deleted refresh tokens (should be 1 of each)
+    async with auth_db as auth_db:
+        deleted_expired, deleted_revoked = await auth_db.clean_expired_refresh_token(
+            max_validity=0, max_retention=0
+        )
+    assert deleted_expired == 1
+    assert deleted_revoked == 1
+
+    # Get all refresh tokens (Admin)
+    async with auth_db as auth_db:
+        refresh_tokens = await auth_db.get_user_refresh_tokens()
+
+    # Check the number of retrieved refresh tokens (should be 0)
+    assert len(refresh_tokens) == 0