Set is_canonical flag based on account email on test environments (#5087)

Author: bartaz

tests/login/tests_login_handler.py

@@ -5,6 +5,8 @@
 from pymacaroons import Macaroon
 from webapp.app import create_app
 
+from unittest.mock import patch, MagicMock
+
 
 class LoginHandlerTest(TestCase):
     def setUp(self):
@@ -96,3 +98,114 @@ def test_login_connection_error(self):
         response = self.client.get(self.endpoint_url)
 
         assert response.status_code == 502
+
+
+class AfterLoginHandlerTest(TestCase):
+    def create_app(self):
+        app = create_app(testing=True)
+
+        # set up a fake route for testing the after_login function
+        # since it is decorated with @open_id.after_login
+        @app.route("/_test_after_login")
+        def _test_after_login():
+            from webapp.login.views import after_login
+
+            return after_login(self.mock_resp)
+
+        return app
+
+    # creates a mocked responses for the get_account function
+    # and the login response passed to after_login
+    def prepare_mock_response(
+        self, mock_get_account, email="test@test.com", groups=[]
+    ):
+        self.mock_resp = MagicMock()
+        self.mock_resp.nickname = "test"
+        self.mock_resp.identity_url = "https://login.ubuntu.com/test"
+        self.mock_resp.fullname = "Test"
+        self.mock_resp.image = "test.png"
+        self.mock_resp.email = email
+        self.mock_resp.extensions = {
+            "macaroon": MagicMock(discharge="some-discharge"),
+            "lp": MagicMock(is_member=groups),
+        }
+
+        mock_get_account.return_value = {
+            "username": self.mock_resp.nickname,
+            "displayname": self.mock_resp.fullname,
+            "email": email,
+            "stores": [],
+        }
+
+    @patch("webapp.login.views.ENVIRONMENT", "staging")
+    @patch("webapp.login.views.dashboard.get_stores", return_value=[])
+    @patch("webapp.login.views.logic.get_stores", return_value=[])
+    @patch(
+        "webapp.login.views.dashboard.get_validation_sets", return_value=None
+    )
+    @patch("webapp.login.views.dashboard.get_account")
+    def test_is_canonical_true_if_email_ends_with_canonical_on_staging(
+        self,
+        mock_get_account,
+        *_,
+    ):
+        # on test environments, we treat publisher's account as "canonical"
+        # if their email is (at)canonical email
+        self.prepare_mock_response(
+            mock_get_account, email="test@canonical.com", groups=[]
+        )
+
+        self.client.get("/_test_after_login")
+
+        with self.client.session_transaction() as s:
+            publisher = s.get("publisher")
+            assert publisher is not None
+            assert publisher["is_canonical"] is True
+
+    @patch("webapp.login.views.ENVIRONMENT", "production")
+    @patch("webapp.login.views.dashboard.get_stores", return_value=[])
+    @patch("webapp.login.views.logic.get_stores", return_value=[])
+    @patch(
+        "webapp.login.views.dashboard.get_validation_sets", return_value=None
+    )
+    @patch("webapp.login.views.dashboard.get_account")
+    def test_is_canonical_true_if_member_of_team_on_production(
+        self,
+        mock_get_account,
+        *_,
+    ):
+        # on production, we treat publisher's account as "canonical"
+        # if they are a member of the canonical team
+        self.prepare_mock_response(mock_get_account, groups=["canonical"])
+
+        self.client.get("/_test_after_login")
+
+        with self.client.session_transaction() as s:
+            publisher = s.get("publisher")
+            assert publisher is not None
+            assert publisher["is_canonical"] is True
+
+    @patch("webapp.login.views.ENVIRONMENT", "production")
+    @patch("webapp.login.views.dashboard.get_stores", return_value=[])
+    @patch("webapp.login.views.logic.get_stores", return_value=[])
+    @patch(
+        "webapp.login.views.dashboard.get_validation_sets", return_value=None
+    )
+    @patch("webapp.login.views.dashboard.get_account")
+    def test_is_canonical_false_if_not_member_of_team_on_production(
+        self,
+        mock_get_account,
+        *_,
+    ):
+        # on production, we treat publisher's account as "canonical"
+        # only if they are a member of the canonical team, not based on email
+        self.prepare_mock_response(
+            mock_get_account, email="test@canonical.com", groups=[]
+        )
+
+        self.client.get("/_test_after_login")
+
+        with self.client.session_transaction() as s:
+            publisher = s.get("publisher")
+            assert publisher is not None
+            assert publisher["is_canonical"] is False

webapp/login/views.py

@@ -20,6 +20,14 @@
 )
 
 LOGIN_URL = os.getenv("LOGIN_URL", "https://login.ubuntu.com")
+ENVIRONMENT = os.getenv("ENVIRONMENT", "devel")
+
+
+# getter for ENVIRONMENT variable
+# this allows the value to be mocked in tests
+def get_environment():
+    return ENVIRONMENT
+
 
 LP_CANONICAL_TEAM = "canonical"
 
@@ -75,14 +83,23 @@ def after_login(resp):
     validation_sets = dashboard.get_validation_sets(flask.session)
 
     if account:
+        is_canonical = LP_CANONICAL_TEAM in resp.extensions["lp"].is_member
+
+        # in environments other than production, for testing purposes,
+        # we detect if the user is Canonical by checking
+        # if the email ends with @canonical.com
+        if (not is_canonical) and get_environment() != "production":
+            is_canonical = account["email"] and account["email"].endswith(
+                "@canonical.com"
+            )
+
         flask.session["publisher"] = {
             "identity_url": resp.identity_url,
             "nickname": account["username"],
             "fullname": account["displayname"],
             "image": resp.image,
             "email": account["email"],
-            "is_canonical": LP_CANONICAL_TEAM
-            in resp.extensions["lp"].is_member,
+            "is_canonical": is_canonical,
         }
 
         if logic.get_stores(