Implement test for x_token_server error bug

trinodb · Sep 13, 2024 · 7b6bd59 · 7b6bd59
1 parent feca584
commit 7b6bd59
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/tests/unit/oauth_test_utils.py b/tests/unit/oauth_test_utils.py
@@ -54,11 +54,14 @@ def __call__(self, request, uri, response_headers):
         if authorization and authorization.replace("Bearer ", "") in self.tokens:
             return [200, response_headers, json.dumps(self.sample_post_response_data)]
         elif self.redirect_server is None and self.token_server is not None:
-            return [401, {'Www-Authenticate': f'Bearer x_token_server="{self.token_server}"',
-                          'Basic realm': '"Trino"'}, ""]
-        return [401, {'Www-Authenticate': f'Bearer x_redirect_server="{self.redirect_server}", '
-                                          f'x_token_server="{self.token_server}"',
-                      'Basic realm': '"Trino"'}, ""]
+            return [401,
+                    {'Www-Authenticate': f'Bearer realm="Trino", token_type="JWT", Bearer x_token_server="{self.token_server}"',
+                     'Basic realm': '"Trino"'},
+                    ""]
+        return [401,
+                {'Www-Authenticate': f'Bearer realm="Trino", token_type="JWT", Bearer x_redirect_server="{self.redirect_server}", x_token_server="{self.token_server}"',
+                 'Basic realm': '"Trino"'},
+                ""]
 
 
 class GetTokenCallback: