Spaces:
Build error
Build error
File size: 6,922 Bytes
0fea559 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 |
import unittest
from unittest.mock import patch
from fastapi import Request
from huggingface_hub.utils import capture_output, is_gradio_available
from .testing_utils import require_webhooks
if is_gradio_available():
import gradio as gr
from fastapi.testclient import TestClient
import huggingface_hub._webhooks_server
from huggingface_hub import WebhookPayload, WebhooksServer
# Taken from https://huggingface.co/docs/hub/webhooks#event
WEBHOOK_PAYLOAD_EXAMPLE = {
"event": {"action": "create", "scope": "discussion"},
"repo": {
"type": "model",
"name": "gpt2",
"id": "621ffdc036468d709f17434d",
"private": False,
"url": {"web": "https://huggingface.co/gpt2", "api": "https://huggingface.co/api/models/gpt2"},
"owner": {"id": "628b753283ef59b5be89e937"},
},
"discussion": {
"id": "6399f58518721fdd27fc9ca9",
"title": "Update co2 emissions",
"url": {
"web": "https://huggingface.co/gpt2/discussions/19",
"api": "https://huggingface.co/api/models/gpt2/discussions/19",
},
"status": "open",
"author": {"id": "61d2f90c3c2083e1c08af22d"},
"num": 19,
"isPullRequest": True,
"changes": {"base": "refs/heads/main"},
},
"comment": {
"id": "6399f58518721fdd27fc9caa",
"author": {"id": "61d2f90c3c2083e1c08af22d"},
"content": "Add co2 emissions information to the model card",
"hidden": False,
"url": {"web": "https://huggingface.co/gpt2/discussions/19#6399f58518721fdd27fc9caa"},
},
"webhook": {"id": "6390e855e30d9209411de93b", "version": 3},
}
@require_webhooks
class TestWebhooksServerDontRun(unittest.TestCase):
def test_add_webhook_implicit_path(self):
# Test adding a webhook
app = WebhooksServer()
@app.add_webhook
async def handler():
pass
self.assertIn("/webhooks/handler", app.registered_webhooks)
def test_add_webhook_explicit_path(self):
# Test adding a webhook
app = WebhooksServer()
@app.add_webhook(path="/test_webhook")
async def handler():
pass
self.assertIn("/webhooks/test_webhook", app.registered_webhooks) # still registered under /webhooks
def test_add_webhook_twice_should_fail(self):
# Test adding a webhook
app = WebhooksServer()
@app.add_webhook("my_webhook")
async def test_webhook():
pass
# Registering twice the same webhook should raise an error
with self.assertRaises(ValueError):
@app.add_webhook("my_webhook")
async def test_webhook_2():
pass
@require_webhooks
class TestWebhooksServerRun(unittest.TestCase):
HEADERS_VALID_SECRET = {"x-webhook-secret": "my_webhook_secret"}
HEADERS_WRONG_SECRET = {"x-webhook-secret": "wrong_webhook_secret"}
def setUp(self) -> None:
with gr.Blocks() as ui:
gr.Markdown("Hello World!")
app = WebhooksServer(ui=ui, webhook_secret="my_webhook_secret")
# Route to check payload parsing
@app.add_webhook
async def test_webhook(payload: WebhookPayload) -> None:
return {"scope": payload.event.scope}
# Routes to check secret validation
# Checks all 4 cases (async/sync, with/without request parameter)
@app.add_webhook
async def async_with_request(request: Request) -> None:
return {"success": True}
@app.add_webhook
def sync_with_request(request: Request) -> None:
return {"success": True}
@app.add_webhook
async def async_no_request() -> None:
return {"success": True}
@app.add_webhook
def sync_no_request() -> None:
return {"success": True}
# Route to check explicit path
@app.add_webhook(path="/explicit_path")
async def with_explicit_path() -> None:
return {"success": True}
self.ui = ui
self.app = app
self.client = self.mocked_run_app()
def tearDown(self) -> None:
self.ui.server.close()
def mocked_run_app(self) -> "TestClient":
with patch.object(self.ui, "block_thread"):
# Run without blocking
with patch.object(huggingface_hub._webhooks_server, "_is_local", False):
# Run without tunnel
self.app.run()
return TestClient(self.app.fastapi_app)
def test_run_print_instructions(self):
"""Test that the instructions are printed when running the app."""
# Test running the app
with capture_output() as output:
self.mocked_run_app()
instructions = output.getvalue()
self.assertIn("Webhooks are correctly setup and ready to use:", instructions)
self.assertIn("- POST http://127.0.0.1:7860/webhooks/test_webhook", instructions)
def test_run_parse_payload(self):
"""Test that the payload is correctly parsed when running the app."""
response = self.client.post(
"/webhooks/test_webhook", headers=self.HEADERS_VALID_SECRET, json=WEBHOOK_PAYLOAD_EXAMPLE
)
self.assertEqual(response.status_code, 200)
self.assertEqual(response.json(), {"scope": "discussion"})
def test_with_webhook_secret_should_succeed(self):
"""Test success if valid secret is sent."""
for path in ["async_with_request", "sync_with_request", "async_no_request", "sync_no_request"]:
with self.subTest(path):
response = self.client.post(f"/webhooks/{path}", headers=self.HEADERS_VALID_SECRET)
self.assertEqual(response.status_code, 200)
self.assertEqual(response.json(), {"success": True})
def test_no_webhook_secret_should_be_unauthorized(self):
"""Test failure if valid secret is sent."""
for path in ["async_with_request", "sync_with_request", "async_no_request", "sync_no_request"]:
with self.subTest(path):
response = self.client.post(f"/webhooks/{path}")
self.assertEqual(response.status_code, 401)
def test_wrong_webhook_secret_should_be_forbidden(self):
"""Test failure if valid secret is sent."""
for path in ["async_with_request", "sync_with_request", "async_no_request", "sync_no_request"]:
with self.subTest(path):
response = self.client.post(f"/webhooks/{path}", headers=self.HEADERS_WRONG_SECRET)
self.assertEqual(response.status_code, 403)
def test_route_with_explicit_path(self):
"""Test that the route with an explicit path is correctly registered."""
response = self.client.post("/webhooks/explicit_path", headers=self.HEADERS_VALID_SECRET)
self.assertEqual(response.status_code, 200)
|