Merge pull request #9196 from khoaguin/bigquery-mock-client

Create mock bigquery client so bigquery API endpoints can be run without a real API

notebooks/scenarios/bigquery/01-setup-datasite.ipynb

@@ -256,16 +256,14 @@
    "source": [
     "# !docker image ls | grep bigquery"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
@@ -276,7 +274,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.5"
+   "version": "3.12.2"
   }
  },
  "nbformat": 4,

notebooks/scenarios/bigquery/02-configure-api.ipynb

@@ -755,16 +755,14 @@
    "source": [
     "server.land()"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
@@ -775,7 +773,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.5"
+   "version": "3.12.2"
   }
  },
  "nbformat": 4,

notebooks/scenarios/bigquery/03-ds-submit-request.ipynb

@@ -223,16 +223,14 @@
    "source": [
     "server.land()"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
@@ -243,7 +241,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.5"
+   "version": "3.12.2"
   }
  },
  "nbformat": 4,

notebooks/scenarios/bigquery/04-do-review-requests.ipynb

@@ -243,16 +243,14 @@
    "source": [
     "server.land()"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
@@ -263,7 +261,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.5"
+   "version": "3.12.2"
   }
  },
  "nbformat": 4,

notebooks/scenarios/bigquery/05-ds-get-results.ipynb

@@ -128,6 +128,11 @@
   }
  ],
  "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
@@ -138,7 +143,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.5"
+   "version": "3.12.2"
   }
  },
  "nbformat": 4,

packages/syft/src/syft/util/util.py

@@ -38,6 +38,7 @@
 import types
 from types import ModuleType
 from typing import Any
+from unittest.mock import Mock
 
 # third party
 from IPython.display import display
@@ -1143,3 +1144,47 @@ def repr_truncation(obj: Any, max_elements: int = 10) -> str:
     r.maxother = 100  # For other objects
 
     return r.repr(obj)
+
+
+class MockBigQueryError(Exception):
+    def __init__(self, errors: list) -> None:
+        self._errors = errors
+        super().__init__(self._errors[0]["message"])
+
+
+class MockBigQueryClient:
+    def __init__(self, credentials: dict, location: str | None = None) -> None:
+        self.credentials = credentials
+        self.location = location
+
+    def query_and_wait(self, sql_query: str, project: str | None = None) -> Mock | None:
+        if self.credentials["mock_result"] == "timeout":
+            raise TimeoutError("Simulated query timeout.")
+
+        if self.credentials["mock_result"] == "success":
+            # Simulate a successful response
+            rows = Mock()
+            rows.total_rows = 1  # or any number within acceptable limits
+            rows.to_dataframe = Mock(return_value="Simulated DataFrame")
+            return rows
+
+        if self.credentials["mock_result"] == "bigquery_error":
+            errors = [
+                {
+                    "reason": "Simulated BigQuery error.",
+                    "message": "Simulated BigQuery error.",
+                }
+            ]
+            raise MockBigQueryError(errors)
+
+        raise Exception("Simulated non-BigQuery exception.")
+
+
+class MockBigQuery:
+    @staticmethod
+    def mock_credentials(mock_result: str = "success") -> dict:
+        return {"mocked": "credentials", "mock_result": mock_result}
+
+    @staticmethod
+    def Client(credentials: dict, location: str | None = None) -> MockBigQueryClient:
+        return MockBigQueryClient(credentials, location)