zulip/zerver/tests/test_subdomains.py

from typing import Any, Mapping, Sequence
from unittest import mock

from zerver.lib.subdomains import get_subdomain
from zerver.lib.test_classes import ZulipTestCase
from zerver.models import Realm


class SubdomainsTest(ZulipTestCase):
    def test_get_subdomain(self) -> None:

        def request_mock(host: str) -> Any:
            request = mock.Mock(spec=['get_host'])
            request.attach_mock(mock.Mock(return_value=host), 'get_host')
            return request

        def test(expected: str, host: str, *, plusport: bool=True,
                 external_host: str='example.org',
                 realm_hosts: Mapping[str, str]={},
                 root_aliases: Sequence[str]=[]) -> None:
            with self.settings(EXTERNAL_HOST=external_host,
                               REALM_HOSTS=realm_hosts,
                               ROOT_SUBDOMAIN_ALIASES=root_aliases):
                self.assertEqual(get_subdomain(request_mock(host)), expected)
                if plusport and ':' not in host:
                    self.assertEqual(get_subdomain(request_mock(host + ':443')),
                                     expected)

        ROOT = Realm.SUBDOMAIN_FOR_ROOT_DOMAIN

        # Basics
        test(ROOT, 'example.org')
        test('foo', 'foo.example.org')
        test(ROOT, 'www.example.org', root_aliases=['www'])

        # Unrecognized patterns fall back to root
        test(ROOT, 'arbitrary.com')
        test(ROOT, 'foo.example.org.evil.com')

        # REALM_HOSTS adds a name,
        test('bar', 'chat.barbar.com', realm_hosts={'bar': 'chat.barbar.com'})
        # ... exactly, ...
        test(ROOT, 'surchat.barbar.com', realm_hosts={'bar': 'chat.barbar.com'})
        test(ROOT, 'foo.chat.barbar.com', realm_hosts={'bar': 'chat.barbar.com'})
        # ... and leaves the subdomain in place too.
        test('bar', 'bar.example.org', realm_hosts={'bar': 'chat.barbar.com'})

        # Any port is fine in Host if there's none in EXTERNAL_HOST, ...
        test('foo', 'foo.example.org:443', external_host='example.org')
        test('foo', 'foo.example.org:12345', external_host='example.org')
        # ... but an explicit port in EXTERNAL_HOST must be explicitly matched in Host.
        test(ROOT, 'foo.example.org', external_host='example.org:12345')
        test(ROOT, 'foo.example.org', external_host='example.org:443', plusport=False)
        test('foo', 'foo.example.org:443', external_host='example.org:443')