From 315eb88900809938401bea855e8f407da1cb509c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebasti=C3=A1n=20Ram=C3=ADrez?= <tiangolo@gmail.com>
Date: Mon, 10 Dec 2018 20:12:28 +0400
Subject: [PATCH] :white_check_mark: Add tests for OAuth2 Password scheme

---
 tests/main.py          | 34 ++++++++++++++++++++++++++++------
 tests/test_security.py | 26 ++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_security.py

diff --git a/tests/main.py b/tests/main.py
index 0c66c1ff9..23f06dff3 100644
--- a/tests/main.py
+++ b/tests/main.py
@@ -1,3 +1,6 @@
+from typing import Optional
+
+import fastapi
 from fastapi import (
     Body,
     Cookie,
@@ -10,8 +13,14 @@ from fastapi import (
     Query,
     Security,
 )
-from fastapi.security import HTTPBasic, OAuth2, OAuth2PasswordRequestForm
+from fastapi.security import (
+    HTTPBasic,
+    OAuth2,
+    OAuth2PasswordBearer,
+    OAuth2PasswordRequestForm,
+)
 from pydantic import BaseModel
+from starlette.exceptions import HTTPException
 from starlette.responses import HTMLResponse, JSONResponse, PlainTextResponse
 from starlette.status import HTTP_202_ACCEPTED
 from starlette.testclient import TestClient
@@ -240,14 +249,27 @@ def get_security_oauth2(sec=Security(reusable_oauth2, scopes=["read:user"])):
     return sec
 
 
+reusable_oauth2b = OAuth2PasswordBearer(tokenUrl="/token")
+
+
+class User(BaseModel):
+    username: str
+
+
+def get_current_user(oauth_header: str = Security(reusable_oauth2b)):
+    user = User(username=oauth_header)
+    return user
+
+
+@app.get("/security/oauth2b")
+def read_current_user(current_user: User = Depends(get_current_user)):
+    return current_user
+
+
 @app.post("/token")
 def post_token(request_data: OAuth2PasswordRequestForm = Form(...)):
-    print(request_data)
     data = request_data.parse()
-    print(data)
-
-    print(request_data())
-    access_token = request_data.username + ":" + request_data.password
+    access_token = data.username + ":" + data.password
     return {"access_token": access_token}
 
 
diff --git a/tests/test_security.py b/tests/test_security.py
new file mode 100644
index 000000000..25ec0c51f
--- /dev/null
+++ b/tests/test_security.py
@@ -0,0 +1,26 @@
+import pytest
+from starlette.testclient import TestClient
+
+from .main import app
+
+client = TestClient(app)
+
+
+def test_security_oauth2_password_bearer():
+    response = client.get(
+        "/security/oauth2b", headers={"Authorization": "Bearer footokenbar"}
+    )
+    assert response.status_code == 200
+    assert response.json() == {"username": "footokenbar"}
+
+
+def test_security_oauth2_password_bearer_wrong_header():
+    response = client.get("/security/oauth2b", headers={"Authorization": "footokenbar"})
+    assert response.status_code == 403
+    assert response.json() == {"detail": "Not authenticated"}
+
+
+def test_security_oauth2_password_bearer_no_header():
+    response = client.get("/security/oauth2b")
+    assert response.status_code == 403
+    assert response.json() == {"detail": "Not authenticated"}