Added tests and adjusted password reset endpoints

app/apis/auth/service.py

@@ -10,10 +10,11 @@
     create_access_token,
     create_user,
     get_current_user,
+    send_code_to_phone_number,
     update_user,
     update_user_password,
 )
-from db.models import User
+from db.models import User, UserRole
 from db.session import get_db
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 from fastapi.security import OAuth2PasswordRequestForm
@@ -113,13 +114,16 @@ def reset_password(
     db: Session = Depends(get_db),
 ):
     user = db.query(User).filter(User.username == data.username).first()
-
-    print(user)
     if not user:
         raise HTTPException(
             status_code=400,
             detail="Invalid username or phone number",
         )
+    if user.role != UserRole.CUSTOMER:
+        raise HTTPException(
+            status_code=400,
+            detail="Only customers can reset their password through this feature",
+        )
 
     if user.phone_number.replace(" ", "") != data.phone_number.replace(" ", ""):
         raise HTTPException(
@@ -134,8 +138,7 @@ def reset_password(
     db.add(user)
     db.commit()
 
-    print(f"PIN {user.reset_password_code} set for {user.username}")
-
+    send_code_to_phone_number(user.phone_number, user.reset_password_code)
     return {"detail": "PIN code sent to your phone number"}
 
 
@@ -148,7 +151,6 @@ def set_new_password(
     db: Session = Depends(get_db),
 ):
     user = db.query(User).filter(User.username == data.username).first()
-
     if not user:
         raise HTTPException(
             status_code=400,
@@ -180,5 +182,9 @@ def set_new_password(
         )
 
     update_user_password(db, user.username, data.new_password)
+    user.reset_password_code = None
+    user.reset_password_code_expiry_date = None
+    db.add(user)
+    db.commit()
 
-    return {}
+    return {"detail": "Password updated successfully!"}

app/apis/auth/utils.py

@@ -166,3 +166,10 @@ def __init__(
     def __call__(self, user: User = Depends(get_current_user)):
         if user.role not in self.required_roles:
             raise HTTPException(status_code=403, detail="Unauthorized")
+
+
+def send_code_to_phone_number(phone_number: str, code: str):
+    # normally this would send a code to the phone number using
+    # a third party service
+    print(f"Sending code {code} to phone number {phone_number}")
+    return True

app/tests/modules/auth/test_auth_service.py

@@ -1,3 +1,6 @@
+from datetime import datetime, timedelta
+
+from apis.auth.utils import verify_password
 from db.models import User, UserRole
 
 
@@ -127,3 +130,114 @@ def test_register_by_authenticated_user_returns_400(test_db, customer_client, mo
         "/register", json=data, headers={"Authorization": "token"}
     )
     assert response.status_code == 400
+
+
+def test_reset_password_sets_reset_code_and_returns_200(test_db, anon_client):
+    user = User(
+        username="customer",
+        password="password",
+        first_name="Customer",
+        last_name="",
+        phone_number="12345678",
+        role=UserRole.CUSTOMER,
+    )
+    test_db.add(user)
+    test_db.commit()
+
+    data = {
+        "username": "customer",
+        "phone_number": "12345678",
+    }
+    response = anon_client.post("/reset-password", json=data)
+
+    user = test_db.query(User).filter(User.username == "customer").first()
+    assert user.reset_password_code is not None
+    assert user.reset_password_code_expiry_date is not None
+    assert response.status_code == 200
+    assert response.json().get("detail") == "PIN code sent to your phone number"
+
+
+def test_reset_password_with_invalid_data_returns_400(test_db, anon_client):
+    data = {
+        "username": "invalid_user",
+        "phone_number": "12345678",
+    }
+    response = anon_client.post("/reset-password", json=data)
+
+    assert response.status_code == 400
+    assert response.json().get("detail") == "Invalid username or phone number"
+
+
+def test_reset_password_for_non_customer_returns_400(test_db, anon_client):
+    user = User(
+        username="employee",
+        password="password",
+        first_name="Employee",
+        last_name="",
+        phone_number="12345678",
+        role=UserRole.EMPLOYEE,
+    )
+    test_db.add(user)
+    test_db.commit()
+
+    data = {
+        "username": "employee",
+        "phone_number": "12345678",
+    }
+    response = anon_client.post("/reset-password", json=data)
+
+    assert response.status_code == 400
+    assert (
+        response.json().get("detail")
+        == "Only customers can reset their password through this feature"
+    )
+
+
+def test_reset_password_with_incorrect_phone_number_returns_400(test_db, anon_client):
+    user = User(
+        username="customer",
+        password="password",
+        first_name="Customer",
+        last_name="",
+        phone_number="12345678",
+        role=UserRole.CUSTOMER,
+    )
+    test_db.add(user)
+    test_db.commit()
+
+    data = {
+        "username": "customer",
+        "phone_number": "12345679",
+    }
+    response = anon_client.post("/reset-password", json=data)
+
+    assert response.status_code == 400
+    assert response.json().get("detail") == "Invalid username or phone number"
+
+
+def test_set_new_password_returns_200(test_db, anon_client):
+    user = User(
+        username="customer",
+        password="password",
+        first_name="Customer",
+        last_name="",
+        phone_number="12345678",
+        role=UserRole.CUSTOMER,
+        reset_password_code="1234",
+        reset_password_code_expiry_date=datetime.now() + timedelta(minutes=15),
+    )
+    test_db.add(user)
+    test_db.commit()
+
+    data = {
+        "username": "customer",
+        "phone_number": "12345678",
+        "reset_password_code": "1234",
+        "new_password": "new_password",
+    }
+    response = anon_client.post("/reset-password/new-password", json=data)
+
+    user = test_db.query(User).filter(User.username == "customer").first()
+    assert verify_password("new_password", user.password)
+    assert response.status_code == 200
+    assert response.json().get("detail") == "Password updated successfully!"