From b14d692ef12ba3e39108b31303e40ac0a8414ea1 Mon Sep 17 00:00:00 2001
From: Robert Jambrecic <robert@airt.ai>
Date: Wed, 19 Jun 2024 14:40:56 +0200
Subject: [PATCH] Refactoring and tests added

---
 google_sheets/app.py  | 36 +++++++++++------------
 tests/app/test_app.py | 67 ++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 84 insertions(+), 19 deletions(-)

diff --git a/google_sheets/app.py b/google_sheets/app.py
index 516fdc9..3bcd57c 100644
--- a/google_sheets/app.py
+++ b/google_sheets/app.py
@@ -97,22 +97,6 @@ def _get_sheet(service: Any, spreadsheet_id: str, range: str) -> Any:
     return values
 
 
-@asyncify  # type: ignore[misc]
-def list_sheets(service: Any) -> List[Dict[str, str]]:
-    # Call the Drive v3 API
-    results = (
-        service.files()
-        .list(
-            q="mimeType='application/vnd.google-apps.spreadsheet'",
-            pageSize=10,
-            fields="nextPageToken, files(id, name)",
-        )
-        .execute()
-    )
-    items = results.get("files", [])
-    return items  # type: ignore[no-any-return]
-
-
 @app.get("/sheet", description="Get data from a Google Sheet")
 async def get_sheet(
     user_id: Annotated[
@@ -137,14 +121,30 @@ async def get_sheet(
     return values  # type: ignore[no-any-return]
 
 
-@app.get("/all", description="Get all sheets associated with the user")
+@asyncify  # type: ignore[misc]
+def _get_files(service: Any) -> List[Dict[str, str]]:
+    # Call the Drive v3 API
+    results = (
+        service.files()
+        .list(
+            q="mimeType='application/vnd.google-apps.spreadsheet'",
+            pageSize=100,  # The default value is 100
+            fields="nextPageToken, files(id, name)",
+        )
+        .execute()
+    )
+    items = results.get("files", [])
+    return items  # type: ignore[no-any-return]
+
+
+@app.get("/get-all-file-names", description="Get all sheets associated with the user")
 async def get_all_file_names(
     user_id: Annotated[
         int, Query(description="The user ID for which the data is requested")
     ],
 ) -> Dict[str, str]:
     service = await _build_service(user_id=user_id, service_name="drive", version="v3")
-    files: List[Dict[str, str]] = await list_sheets(service=service)
+    files: List[Dict[str, str]] = await _get_files(service=service)
     # create dict where key is id and value is name
     files_dict = {file["id"]: file["name"] for file in files}
     return files_dict
diff --git a/tests/app/test_app.py b/tests/app/test_app.py
index 4f541c7..af8e8ba 100644
--- a/tests/app/test_app.py
+++ b/tests/app/test_app.py
@@ -31,6 +31,27 @@ def test_get_sheet(self) -> None:
                 assert response.status_code == 200
                 assert response.json() == excepted
 
+    def test_get_all_file_names(self) -> None:
+        with (
+            patch(
+                "google_sheets.app.load_user_credentials",
+                return_value={"refresh_token": "abcdf"},
+            ) as mock_load_user_credentials,
+            patch(
+                "google_sheets.app._get_files",
+                return_value=[
+                    {"id": "abc", "name": "file1"},
+                    {"id": "def", "name": "file2"},
+                ],
+            ) as mock_get_files,
+        ):
+            expected = {"abc": "file1", "def": "file2"}
+            response = client.get("/get-all-file-names?user_id=123")
+            mock_load_user_credentials.assert_called_once()
+            mock_get_files.assert_called_once()
+            assert response.status_code == 200
+            assert response.json() == expected
+
     def test_openapi(self) -> None:
         expected = {
             "openapi": "3.1.0",
@@ -115,7 +136,51 @@ def test_openapi(self) -> None:
                             },
                         },
                     }
-                }
+                },
+                "/get-all-file-names": {
+                    "get": {
+                        "summary": "Get All File Names",
+                        "description": "Get all sheets associated with the user",
+                        "operationId": "get_all_file_names_get_all_file_names_get",
+                        "parameters": [
+                            {
+                                "name": "user_id",
+                                "in": "query",
+                                "required": True,
+                                "schema": {
+                                    "type": "integer",
+                                    "description": "The user ID for which the data is requested",
+                                    "title": "User Id",
+                                },
+                                "description": "The user ID for which the data is requested",
+                            }
+                        ],
+                        "responses": {
+                            "200": {
+                                "description": "Successful Response",
+                                "content": {
+                                    "application/json": {
+                                        "schema": {
+                                            "type": "object",
+                                            "additionalProperties": {"type": "string"},
+                                            "title": "Response Get All File Names Get All File Names Get",
+                                        }
+                                    }
+                                },
+                            },
+                            "422": {
+                                "description": "Validation Error",
+                                "content": {
+                                    "application/json": {
+                                        "schema": {
+                                            "$ref": "#/components/schemas/HTTPValidationError"
+                                        }
+                                    }
+                                },
+                            },
+                        },
+                    }
+                },
             },
             "components": {
                 "schemas": {