authentik/authentik/flows/tests/__init__.py

"""Test helpers"""

from json import loads
from typing import Any

from django.http.response import HttpResponse
from django.urls.base import reverse
from rest_framework.test import APITestCase

from authentik.core.models import User
from authentik.flows.challenge import ChallengeTypes
from authentik.flows.models import Flow


class FlowTestCase(APITestCase):
    """Helpers for testing flows and stages."""

    def assertStageResponse(
        self,
        response: HttpResponse,
        flow: Flow | None = None,
        user: User | None = None,
        **kwargs,
    ) -> dict[str, Any]:
        """Assert various attributes of a stage response"""
        self.assertEqual(response.status_code, 200)
        raw_response = loads(response.content.decode())
        self.assertIsNotNone(raw_response["component"])
        self.assertIsNotNone(raw_response["type"])
        if flow:
            self.assertIn("flow_info", raw_response)
            self.assertEqual(raw_response["flow_info"]["background"], flow.background_url)
            self.assertEqual(
                raw_response["flow_info"]["cancel_url"], reverse("authentik_flows:cancel")
            )
            # We don't check the flow title since it will most likely go
            # through ChallengeStageView.format_title() so might not match 1:1
            # self.assertEqual(raw_response["flow_info"]["title"], flow.title)
            self.assertIsNotNone(raw_response["flow_info"]["title"])
        if user:
            self.assertEqual(raw_response["pending_user"], user.username)
            self.assertEqual(raw_response["pending_user_avatar"], user.avatar)
        for key, expected in kwargs.items():
            self.assertEqual(raw_response[key], expected)
        return raw_response

    def assertStageRedirects(self, response: HttpResponse, to: str) -> dict[str, Any]:
        """Wrapper around assertStageResponse that checks for a redirect"""
        return self.assertStageResponse(
            response, component="xak-flow-redirect", to=to, type=ChallengeTypes.REDIRECT.value
        )