From 082d6d934c9e44e6b41b1f8aac19138da061cfde Mon Sep 17 00:00:00 2001
From: gauravr <gaurav@stck.me>
Date: Wed, 6 Mar 2024 15:48:53 +0530
Subject: [PATCH] Added piccolo+fastapi test

---
 fastapi_tests/app/endpoints.py | 15 +++++++++++++++
 fastapi_tests/service.py       |  2 ++
 fastapi_tests/test_rest.py     | 28 ++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+)

diff --git a/fastapi_tests/app/endpoints.py b/fastapi_tests/app/endpoints.py
index eb3893b..c1e97e9 100644
--- a/fastapi_tests/app/endpoints.py
+++ b/fastapi_tests/app/endpoints.py
@@ -5,6 +5,7 @@
 
 from apphelpers.rest import endpoint as ep
 from apphelpers.rest.fastapi import json_body, user, user_agent, user_id
+from fastapi_tests.app.models import Book
 
 
 def echo(word, user=user):
@@ -93,6 +94,18 @@ async def get_fields(fields: set = Query(..., default_factory=set)):
     return {k: v for k, v in data.items() if k in fields}
 
 
+async def add_books(succeed: bool):
+    await Book.insert(Book(name="The Pillars of the Earth")).run()
+    await Book.insert(Book(name="The Cathedral and the Bazaar")).run()
+    if not succeed:
+        raise ValueError("Failure")
+    await Book.insert(Book(name="The Ego Trick")).run()
+
+
+async def count_books():
+    return await Book.count()
+
+
 def setup_routes(factory):
     factory.get("/echo/{word}")(echo)
     factory.get("/echo-async/{word}")(echo_async)
@@ -117,3 +130,5 @@ def setup_routes(factory):
         echo_user_agent_without_site_ctx_async
     )
     factory.get("/fields")(get_fields)
+    factory.get("/count-books")(count_books)
+    factory.post("/add-books")(add_books)
diff --git a/fastapi_tests/service.py b/fastapi_tests/service.py
index e65b867..6468c3b 100644
--- a/fastapi_tests/service.py
+++ b/fastapi_tests/service.py
@@ -9,6 +9,7 @@
 
 from apphelpers.rest.fastapi import APIFactory
 from fastapi_tests.app.endpoints import setup_routes
+from fastapi_tests.app.models import db
 
 
 def make_app():
@@ -20,6 +21,7 @@ def make_app():
     )
 
     api_factory = APIFactory(sessiondb_conn=sessiondb_conn, site_identifier="site_id")
+    api_factory.setup_db_transaction(db)
     setup_routes(api_factory)
 
     app = fastapi.FastAPI()
diff --git a/fastapi_tests/test_rest.py b/fastapi_tests/test_rest.py
index 09923c1..1f694a8 100644
--- a/fastapi_tests/test_rest.py
+++ b/fastapi_tests/test_rest.py
@@ -2,6 +2,8 @@
 from converge import settings
 
 import apphelpers.sessions as sessionslib
+from apphelpers.db.piccolo import setup_db_from_basetable, destroy_db_from_basetable
+from fastapi_tests.app.models import BaseTable
 
 base_url = "http://127.0.0.1:5000/"
 echo_url = base_url + "echo"
@@ -24,6 +26,13 @@
 
 def setup_module():
     sessionsdb.destroy_all()
+    destroy_db_from_basetable(BaseTable)
+    setup_db_from_basetable(BaseTable)
+
+
+def teardown_module():
+    sessionsdb.destroy_all()
+    destroy_db_from_basetable(BaseTable)
 
 
 def test_get():
@@ -284,3 +293,22 @@ def test_user_agent_async_and_site_ctx():
     response = requests.get(url, headers=headers)
     assert response.status_code == 200
     assert "python-requests" in response.text
+
+
+def test_piccolo():
+    url = base_url + "count-books"
+    assert requests.get(url).json() == 0
+
+    url = base_url + "add-books"
+    data = {"succeed": True}
+    assert requests.post(url, params=data).status_code == 200
+
+    url = base_url + "count-books"
+    assert requests.get(url).json() == 3
+
+    url = base_url + "add-books"
+    data = {"succeed": False}
+    assert requests.post(url, params=data).status_code == 500
+
+    url = base_url + "count-books"
+    assert requests.get(url).json() == 3