fix flaky test for oauth redirects (TNL-4190) (#12187)

common/test/acceptance/tests/lms/test_oauth2.py

@@ -66,7 +66,6 @@ class OAuth2PermissionDelegationTests(WebAppTest):
         else:
             self.oauth_page.wait_for(check_redirect, 'redirected to invalid URL')
 
-    @flaky      # TODO, fix this: TNL-4190
     def test_accepting_redirects(self):
         """
         If you accept the request, you're redirected to the redirect_url with
@@ -77,14 +76,34 @@ class OAuth2PermissionDelegationTests(WebAppTest):
 
         # This redirects to an invalid URI.
         self.oauth_page.confirm()
-        self.oauth_page.wait_for_element_absence('input[name=authorize]', 'Authorization button is not present')
+        self.oauth_page.wait_for_element_absence(
+            'input[name=authorize]', 'Authorization button is not present'
+        )
 
-        # Due to a bug in ChromeDriver, when chrome is on invalid URI,self.browser.current_url outputs
-        # data:text/html,chromewebdata. When this happens in our case,query string is present in the title.
-        # So to get query string, we branch out based on selected browser.
-        if self.browser.name == 'chrome':
-            query = self._qs(self.browser.title)
-        else:
+        def check_query_string():
+            """
+            Checks that 'code' appears in the browser's current url.
+            """
             query = self._qs(self.browser.current_url)
+            return 'code' in query
+
+        def check_query_string_chrome():
+            """
+            Similar to check_query_string, but, due to a bug in ChromeDriver,
+            when chrome is on an invalid URI, `self.browser.current_url` outputs
+            "data:text/html,chromewebdata" instead of the current URI.
+
+            However, since the query string is present in the `title`, we use
+            that for chrome.
+            """
+            query = self._qs(self.browser.title)
+            return 'code' in query
 
-        self.assertIn('code', query)
+        if self.browser.name == 'chrome':
+            self.oauth_page.wait_for(
+                check_query_string_chrome, 'redirected with correct query parameters (chrome)'
+            )
+        else:
+            self.oauth_page.wait_for(
+                check_query_string, 'redirected with correct query parameters'
+            )