Add succes login endpoint

airtai · Jun 20, 2024 · a3bed6a · a3bed6a
1 parent e351e15
commit a3bed6a
Show file tree

Hide file tree

Showing 2 changed files with 127 additions and 5 deletions.
diff --git a/google_sheets/app.py b/google_sheets/app.py
@@ -95,14 +95,18 @@ async def get_login_url(
     return {"login_url": markdown_url}
 
 
+@app.get("/login/success")
+async def get_login_success() -> Dict[str, str]:
+    return {"login_success": "You have successfully logged in"}
+
+
 # Route 2: Save user credentials/token to a JSON file
 @app.get("/login/callback")
 async def login_callback(
     code: str = Query(title="Authorization Code"), state: str = Query(title="State")
 ) -> RedirectResponse:
     chat_id = state
     user_id, chat_uuid = await get_user_id_chat_uuid_from_chat_id(chat_id)
-    # user_id, chat_id = await get_user_id_chat_id_from_conversation(conv_id)
     user = await get_user(user_id=user_id)
 
     token_request_data = {
@@ -145,10 +149,12 @@ async def login_callback(
             },
         )
 
-    redirect_domain = environ.get("REDIRECT_DOMAIN", "https://captn.ai")
-    logged_in_message = "I have successfully logged in"
-    redirect_uri = f"{redirect_domain}/chat/{chat_uuid}?msg={logged_in_message}"
-    return RedirectResponse(redirect_uri)
+    # redirect_domain = environ.get("REDIRECT_DOMAIN", "https://captn.ai")
+    # logged_in_message = "I have successfully logged in"
+    # redirect_uri = f"{redirect_domain}/chat/{chat_uuid}?msg={logged_in_message}"
+    # return RedirectResponse(redirect_uri)
+    # redirect to success page
+    return RedirectResponse(url=f"{base_url}/login/success")
 
 
 async def get_user(user_id: Union[int, str]) -> Any:

diff --git a/tests/app/test_app.py b/tests/app/test_app.py
@@ -63,6 +63,122 @@ def test_openapi(self) -> None:
                 }
             ],
             "paths": {
+                "/login": {
+                    "get": {
+                        "summary": "Get Login Url",
+                        "operationId": "get_login_url_login_get",
+                        "parameters": [
+                            {
+                                "name": "user_id",
+                                "in": "query",
+                                "required": True,
+                                "schema": {"type": "integer", "title": "User ID"},
+                            },
+                            {
+                                "name": "conv_id",
+                                "in": "query",
+                                "required": True,
+                                "schema": {
+                                    "type": "integer",
+                                    "title": "Conversation ID",
+                                },
+                            },
+                            {
+                                "name": "force_new_login",
+                                "in": "query",
+                                "required": False,
+                                "schema": {
+                                    "type": "boolean",
+                                    "title": "Force new login",
+                                    "default": False,
+                                },
+                            },
+                        ],
+                        "responses": {
+                            "200": {
+                                "description": "Successful Response",
+                                "content": {
+                                    "application/json": {
+                                        "schema": {
+                                            "type": "object",
+                                            "additionalProperties": {"type": "string"},
+                                            "title": "Response Get Login Url Login Get",
+                                        }
+                                    }
+                                },
+                            },
+                            "422": {
+                                "description": "Validation Error",
+                                "content": {
+                                    "application/json": {
+                                        "schema": {
+                                            "$ref": "#/components/schemas/HTTPValidationError"
+                                        }
+                                    }
+                                },
+                            },
+                        },
+                    }
+                },
+                "/login/success": {
+                    "get": {
+                        "summary": "Get Login Success",
+                        "operationId": "get_login_success_login_success_get",
+                        "responses": {
+                            "200": {
+                                "description": "Successful Response",
+                                "content": {
+                                    "application/json": {
+                                        "schema": {
+                                            "additionalProperties": {"type": "string"},
+                                            "type": "object",
+                                            "title": "Response Get Login Success Login Success Get",
+                                        }
+                                    }
+                                },
+                            }
+                        },
+                    }
+                },
+                "/login/callback": {
+                    "get": {
+                        "summary": "Login Callback",
+                        "operationId": "login_callback_login_callback_get",
+                        "parameters": [
+                            {
+                                "name": "code",
+                                "in": "query",
+                                "required": True,
+                                "schema": {
+                                    "type": "string",
+                                    "title": "Authorization Code",
+                                },
+                            },
+                            {
+                                "name": "state",
+                                "in": "query",
+                                "required": True,
+                                "schema": {"type": "string", "title": "State"},
+                            },
+                        ],
+                        "responses": {
+                            "200": {
+                                "description": "Successful Response",
+                                "content": {"application/json": {"schema": {}}},
+                            },
+                            "422": {
+                                "description": "Validation Error",
+                                "content": {
+                                    "application/json": {
+                                        "schema": {
+                                            "$ref": "#/components/schemas/HTTPValidationError"
+                                        }
+                                    }
+                                },
+                            },
+                        },
+                    }
+                },
                 "/sheet": {
                     "get": {
                         "summary": "Get Sheet",