rusty-bever/tests/admin.py

65 lines
1.7 KiB
Python
Raw Normal View History

import requests
class RbClient:
def __init__(self, username, password, base_url = "http://localhost:8000/api"):
self.username = username
self.password = password
self.base_url = base_url
self.jwt = None
self.refresh_token = None
def _login(self):
r = requests.post(f"{self.base_url}/auth/login", json={
"username": self.username,
"password": self.password,
})
if r.status_code != 200:
raise Exception("Couldn't login")
res = r.json()
self.jwt = res["token"]
self.refresh_token = res["refreshToken"]
def _refresh(self):
r = requests.post(f"{self.base_url}/auth/refresh", json={"refreshToken": self.refresh_token})
if r.status_code != 200:
raise Exception("Couldn't refresh")
res = r.json()
self.jwt = res["token"]
self.refresh_token = res["refreshToken"]
def _request(self, type_, url, retry=2, *args, **kwargs):
if self.jwt:
headers = kwargs.get("headers", {})
headers["Authorization"] = f"Bearer {self.jwt}"
kwargs["headers"] = headers
print(kwargs["headers"])
r = requests.request(type_, url, *args, **kwargs)
if r.status_code != 200 and retry > 0:
if self.refresh_token:
self._refresh()
else:
self._login()
r = self._request(type_, url, *args, **kwargs, retry=retry - 1)
return r
def get(self, url, *args, **kwargs):
return self._request("GET", f"{self.base_url}{url}", *args, **kwargs)
if __name__ == "__main__":
client = RbClient("admin", "password")
print(client.get("/admin/users").json())