remote_billing: Redirect to upgrade/sponsorship page based on next.

We pass `next` parameter with /self-hosted-billing to redirect
users to the intended page after login.

Fixed realm_uuid incorrectly required in remote_realm_upgrade_page.

corporate/lib/remote_billing_util.py

@@ -17,6 +17,8 @@ class RemoteBillingIdentityDict(TypedDict):
     remote_server_uuid: str
     remote_realm_uuid: str
 
+    next_page: Optional[str]
+
 
 class LegacyServerIdentityDict(TypedDict):
     # Currently this has only one field. We can extend this

corporate/tests/test_remote_billing.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 from unittest import mock
 
 import responses
@@ -16,8 +16,13 @@ if TYPE_CHECKING:
 
 @override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com")
 class RemoteBillingAuthenticationTest(BouncerTestCase):
-    def execute_remote_billing_authentication_flow(self, user: UserProfile) -> "TestHttpResponse":
-        result = self.client_get("/self-hosted-billing/")
+    def execute_remote_billing_authentication_flow(
+        self, user: UserProfile, next_page: Optional[str] = None
+    ) -> "TestHttpResponse":
+        self_hosted_billing_url = "/self-hosted-billing/"
+        if next_page is not None:
+            self_hosted_billing_url += f"?next_page={next_page}"
+        result = self.client_get(self_hosted_billing_url)
         self.assertEqual(result.status_code, 302)
         self.assertIn("http://selfhosting.testserver/remote-billing-login/", result["Location"])
 
@@ -34,6 +39,7 @@ class RemoteBillingAuthenticationTest(BouncerTestCase):
             user_full_name=user.full_name,
             remote_server_uuid=str(self.server.uuid),
             remote_realm_uuid=str(user.realm.uuid),
+            next_page=next_page,
         )
         self.assertEqual(
             self.client.session["remote_billing_identities"][str(user.realm.uuid)], identity_dict
@@ -99,3 +105,43 @@ class RemoteBillingAuthenticationTest(BouncerTestCase):
         result = self.client_get(result["Location"], subdomain="selfhosting")
         self.assert_in_success_response(["Your remote user info:"], result)
         self.assert_in_success_response([desdemona.delivery_email], result)
+
+    @responses.activate
+    def test_remote_billing_authentication_flow_to_sponsorship_page(self) -> None:
+        self.login("desdemona")
+        desdemona = self.example_user("desdemona")
+        realm = desdemona.realm
+
+        self.add_mock_response()
+        send_realms_only_to_push_bouncer()
+
+        result = self.execute_remote_billing_authentication_flow(desdemona, "sponsorship")
+
+        self.assertEqual(result["Location"], f"/realm/{realm.uuid!s}/sponsorship")
+
+        # Go to the URL we're redirected to after authentication and assert
+        # some basic expected content.
+        result = self.client_get(result["Location"], subdomain="selfhosting")
+        self.assert_in_success_response(
+            ["Request Zulip Cloud sponsorship", "Description of your organization"], result
+        )
+
+    @responses.activate
+    def test_remote_billing_authentication_flow_to_upgrade_page(self) -> None:
+        self.login("desdemona")
+        desdemona = self.example_user("desdemona")
+        realm = desdemona.realm
+
+        self.add_mock_response()
+        send_realms_only_to_push_bouncer()
+
+        result = self.execute_remote_billing_authentication_flow(desdemona, "upgrade")
+
+        self.assertEqual(result["Location"], f"/realm/{realm.uuid!s}/upgrade")
+
+        # Go to the URL we're redirected to after authentication and assert
+        # some basic expected content.
+        result = self.client_get(result["Location"], subdomain="selfhosting")
+        self.assert_in_success_response(
+            ["Upgrade", "Purchase Zulip", "Your subscription will renew automatically."], result
+        )

corporate/views/remote_billing_page.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Optional
+from typing import Literal, Optional
 
 from django.conf import settings
 from django.core import signing
@@ -26,6 +26,10 @@ from zilencer.models import RemoteRealm, RemoteZulipServer, get_remote_server_by
 billing_logger = logging.getLogger("corporate.stripe")
 
 
+VALID_NEXT_PAGES = [None, "sponsorship", "upgrade", "billing", "plans"]
+VALID_NEXT_PAGES_TYPE = Literal[None, "sponsorship", "upgrade", "billing", "plans"]
+
+
 @csrf_exempt
 @typed_endpoint
 def remote_server_billing_entry(
@@ -34,6 +38,7 @@ def remote_server_billing_entry(
     *,
     user: Json[UserDataForRemoteBilling],
     realm: Json[RealmDataForAnalytics],
+    next_page: VALID_NEXT_PAGES_TYPE = None,
 ) -> HttpResponse:
     if not settings.DEVELOPMENT:
         return render(request, "404.html", status=404)
@@ -51,6 +56,7 @@ def remote_server_billing_entry(
         user_full_name=user.full_name,
         remote_server_uuid=str(remote_server.uuid),
         remote_realm_uuid=str(remote_realm.uuid),
+        next_page=next_page,
     )
 
     signed_identity_dict = signing.dumps(identity_dict)
@@ -64,7 +70,8 @@ def remote_server_billing_entry(
 
 @self_hosting_management_endpoint
 def remote_server_billing_finalize_login(
-    request: HttpRequest, signed_billing_access_token: str
+    request: HttpRequest,
+    signed_billing_access_token: str,
 ) -> HttpResponse:
     try:
         identity_dict: RemoteBillingIdentityDict = signing.loads(
@@ -85,7 +92,15 @@ def remote_server_billing_finalize_login(
     # For now we're only implemented the case where we have the RemoteRealm, and we take
     # to /plans.
 
-    return HttpResponseRedirect(reverse("remote_billing_plans_realm", args=(remote_realm_uuid,)))
+    assert identity_dict["next_page"] in VALID_NEXT_PAGES
+    if identity_dict["next_page"] is None:
+        return HttpResponseRedirect(
+            reverse("remote_billing_plans_realm", args=(remote_realm_uuid,))
+        )
+    else:
+        return HttpResponseRedirect(
+            reverse(f"remote_realm_{identity_dict['next_page']}_page", args=(remote_realm_uuid,))
+        )
 
 
 def render_tmp_remote_billing_page(

corporate/views/upgrade.py

@@ -21,7 +21,7 @@ from corporate.models import CustomerPlan
 from zerver.decorator import require_organization_member, zulip_login_required
 from zerver.lib.request import REQ, has_request_variables
 from zerver.lib.response import json_success
-from zerver.lib.typed_endpoint import PathOnly, typed_endpoint
+from zerver.lib.typed_endpoint import typed_endpoint
 from zerver.lib.validator import check_bool, check_int, check_string_in
 from zerver.models import UserProfile
 from zilencer.models import RemoteRealm
@@ -108,7 +108,6 @@ def remote_realm_upgrade_page(
     request: HttpRequest,
     remote_realm: RemoteRealm,
     *,
-    realm_uuid: PathOnly[str],
     manual_license_management: Json[bool] = False,
 ) -> HttpResponse:  # nocoverage
     initial_upgrade_request = InitialUpgradeRequest(

web/templates/gear_menu_popover.hbs

@@ -50,6 +50,17 @@
                     {{/unless}}
                     {{/if}}
                 {{/if}}
+
+                {{#if (and is_development_environment show_remote_billing) }}
+                {{! This is only shown in development environment until the UI is ready. Added to help test the flow directly from gear menu.}}
+                {{! TODO: Add these above with proper conditions once the UI is ready. }}
+                <li class="org-upgrade navbar-dropdown-menu-inner-list-item">
+                    <a href="/self-hosted-billing/?next_page=sponsorship" target="_blank" rel="noopener noreferrer" class="navigate-link-on-enter navbar-dropdown-menu-link">{{t "Request sponsorship" }}</a>
+                </li>
+                <li class="org-upgrade navbar-dropdown-menu-inner-list-item">
+                    <a href="/self-hosted-billing/?next_page=upgrade" target="_blank" rel="noopener noreferrer" class="navigate-link-on-enter navbar-dropdown-menu-link">{{t "Upgrade" }}</a>
+                </li>
+                {{/if}}
             </ul>
         </li>
         <li class="navbar-dropdown-menu-outer-list-item">

zerver/views/push_notifications.py

@@ -28,6 +28,7 @@ from zerver.lib.remote_server import (
 )
 from zerver.lib.request import REQ, has_request_variables
 from zerver.lib.response import json_success
+from zerver.lib.typed_endpoint import typed_endpoint
 from zerver.lib.validator import check_string
 from zerver.models import PushDeviceToken, UserProfile
 
@@ -117,7 +118,12 @@ def send_test_push_notification_api(
 
 
 @zulip_login_required
-def self_hosting_auth_redirect(request: HttpRequest) -> HttpResponse:  # nocoverage
+@typed_endpoint
+def self_hosting_auth_redirect(
+    request: HttpRequest,
+    *,
+    next_page: Optional[str] = None,
+) -> HttpResponse:  # nocoverage
     if not settings.DEVELOPMENT or not uses_notification_bouncer():
         return render(request, "404.html", status=404)
 
@@ -143,6 +149,9 @@ def self_hosting_auth_redirect(request: HttpRequest) -> HttpResponse:  # nocover
         "user": user_info.model_dump_json(),
         "realm": realm_info.model_dump_json(),
     }
+    if next_page is not None:
+        post_data["next_page"] = next_page
+
     try:
         result = send_to_push_bouncer("POST", "server/billing", post_data)
     except MissingRemoteRealmError: