Local API (#982)

* Introduce API scelleton * Raw impl for session * Simple state endpoint * Apply _body naming * Finalize session testing and comment out tons of useless code * Clean and refactor part1 * Clean and refactor part2 * Clean and refactor part2 * Clean and refactor part2 * Clean and refactor part2 * Refactor middleware * Refactor middleware * Clean and refactor part3 * TDD and 2 beers * TDD and 2 beers * Complete state endpoints * You can only set what is already present * Use only JSON as response * Use default logger * Improve auth extraction * Extend api command with more args * Adjust API messages
2026-05-12 18:32:18 +02:00 · 2025-11-17 08:36:31 +01:00
parent f1149ebc84
commit 87678da498
20 changed files with 647 additions and 79 deletions
@@ -0,0 +1,154 @@
+import os
+from datetime import timedelta
+from typing import Iterator
+
+import pytest
+from starlette.testclient import TestClient
+
+from facefusion import metadata, session_manager
+from facefusion.apis.core import create_api
+from facefusion.types import Session
+
+
+@pytest.fixture(scope = 'module')
+def test_client() -> Iterator[TestClient]:
+	with TestClient(create_api()) as test_client:
+		yield test_client
+
+
+@pytest.fixture(scope = 'function', autouse = True)
+def before_each() -> None:
+	session_manager.SESSIONS.clear()
+
+
+def test_create_session(test_client : TestClient) -> None:
+	create_session_response = test_client.post('/session', json =
+	{
+		'client_version': metadata.get('version')
+	})
+	create_session_body = create_session_response.json()
+
+	assert session_manager.get_session(create_session_body.get('access_token'))
+	assert create_session_response.status_code == 201
+
+	create_session_response = test_client.post('/session', json =
+	{
+		'api_key': 'TEST',
+		'client_version': metadata.get('version')
+	})
+
+	assert create_session_response.status_code == 401
+
+	os.environ['FACEFUSION_API_KEY'] = 'TEST'
+	create_session_response = test_client.post('/session', json =
+	{
+		'api_key': 'INVALID',
+		'client_version': metadata.get('version')
+	})
+
+	assert create_session_response.status_code == 401
+
+	os.environ['FACEFUSION_API_KEY'] = 'TEST'
+	create_session_response = test_client.post('/session', json =
+	{
+		'api_key': 'TEST',
+		'client_version': metadata.get('version')
+	})
+
+	assert create_session_response.status_code == 201
+
+	del os.environ['FACEFUSION_API_KEY']
+
+
+def test_get_session(test_client : TestClient) -> None:
+	get_session_response = test_client.get('/session')
+
+	assert get_session_response.status_code == 401
+
+	create_session_response = test_client.post('/session', json =
+	{
+		'client_version': metadata.get('version')
+	})
+	create_session_body = create_session_response.json()
+
+	get_session_response = test_client.get('/session', headers =
+	{
+		'Authorization': 'Bearer ' + create_session_body.get('access_token')
+	})
+
+	assert get_session_response.status_code == 200
+
+	access_token = create_session_body.get('access_token')
+	session : Session = session_manager.get_session(access_token)
+	session_manager.set_session(access_token,
+	{
+		'access_token': session.get('access_token'),
+		'refresh_token': session.get('refresh_token'),
+		'created_at': session.get('created_at'),
+		'expires_at': session.get('expires_at') - timedelta(hours = 1)
+	})
+
+	get_session_response = test_client.get('/session', headers =
+	{
+		'Authorization': 'Bearer ' + access_token
+	})
+
+	assert get_session_response.status_code == 426
+
+
+def test_refresh_session(test_client : TestClient) -> None:
+	create_session_response = test_client.post('/session', json =
+	{
+		'client_version': metadata.get('version')
+	})
+	create_session_body = create_session_response.json()
+
+	refresh_session_response = test_client.put('/session', json =
+	{
+		'refresh_token': 'INVALID'
+	})
+
+	assert refresh_session_response.status_code == 401
+
+	refresh_session_response = test_client.put('/session', json =
+	{
+		'refresh_token': create_session_body.get('refresh_token')
+	})
+	refresh_session_body = refresh_session_response.json()
+
+	assert session_manager.get_session(create_session_body.get('access_token')) is None
+
+	assert session_manager.get_session(refresh_session_body.get('access_token'))
+
+	assert refresh_session_response.status_code == 200
+
+	refresh_session_response = test_client.put('/session', json =
+	{
+		'refresh_token': create_session_body.get('refresh_token')
+	})
+
+	assert refresh_session_response.status_code == 401
+
+
+def test_destroy_session(test_client : TestClient) -> None:
+	create_session_response = test_client.post('/session', json =
+	{
+		'client_version': metadata.get('version')
+	})
+	create_session_body = create_session_response.json()
+
+	delete_session_response = test_client.delete('/session', headers =
+	{
+		'Authorization': 'Bearer INVALID'
+	})
+
+	assert delete_session_response.status_code == 401
+
+	delete_session_response = test_client.delete('/session', headers =
+	{
+		'Authorization': 'Bearer ' + create_session_body.get('access_token')
+	})
+
+	assert session_manager.get_session(create_session_body.get('access_token')) is None
+
+	assert delete_session_response.status_code == 200
@@ -0,0 +1,80 @@
+from typing import Iterator
+
+import pytest
+from starlette.testclient import TestClient
+
+from facefusion import metadata, session_manager, state_manager
+from facefusion.apis.core import create_api
+
+
+@pytest.fixture(scope = 'module')
+def test_client() -> Iterator[TestClient]:
+	state_manager.init_item('execution_providers', [ 'cpu' ])
+
+	with TestClient(create_api()) as test_client:
+		yield test_client
+
+
+@pytest.fixture(scope = 'function', autouse = True)
+def before_each() -> None:
+	session_manager.SESSIONS.clear()
+
+
+def test_get_state(test_client : TestClient) -> None:
+	get_state_response = test_client.get('/state')
+
+	assert get_state_response.status_code == 401
+
+	create_session_response = test_client.post('/session', json =
+	{
+		'client_version': metadata.get('version')
+	})
+	create_session_body = create_session_response.json()
+
+	get_state_response = test_client.get('/state', headers =
+	{
+		'Authorization': 'Bearer ' + create_session_body.get('access_token')
+	})
+	get_state_body = get_state_response.json()
+
+	assert get_state_body.get('execution_providers') == [ 'cpu' ]
+	assert get_state_response.status_code == 200
+
+
+def test_set_state(test_client : TestClient) -> None:
+	set_state_response = test_client.put('/state', json =
+	{
+		'execution_providers': [ 'cuda' ]
+	})
+
+	assert set_state_response.status_code == 401
+
+	create_session_response = test_client.post('/session', json =
+	{
+		'client_version': metadata.get('version')
+	})
+	create_session_body = create_session_response.json()
+
+	set_state_response = test_client.put('/state', json =
+	{
+		'execution_providers': [ 'cuda' ]
+	}, headers =
+	{
+		'Authorization': 'Bearer ' + create_session_body.get('access_token')
+	})
+	set_state_body = set_state_response.json()
+
+	assert set_state_body.get('execution_providers') == [ 'cuda' ]
+	assert set_state_response.status_code == 200
+
+	set_state_response = test_client.put('/state', json =
+	{
+		'invalid': 'invalid'
+	}, headers =
+	{
+		'Authorization': 'Bearer ' + create_session_body.get('access_token')
+	})
+	set_state_body = set_state_response.json()
+
+	assert set_state_body.get('invalid') is None
+	assert set_state_response.status_code == 200
@@ -106,7 +106,7 @@ def test_find_jobs() -> None:

 	assert 'job-test-find-jobs-1' in find_jobs('drafted')
 	assert 'job-test-find-jobs-2' in find_jobs('drafted')
-	assert not find_jobs('queued')
+	assert find_jobs('queued') == {}

 	move_job_file('job-test-find-jobs-1', 'queued')

@@ -8,7 +8,7 @@ def test_read_json() -> None:
 	file_descriptor, json_path = tempfile.mkstemp(suffix = '.json')
 	os.close(file_descriptor)

-	assert not read_json(json_path)
+	assert read_json(json_path) is None

 	write_json(json_path, {})

@@ -0,0 +1,45 @@
+import secrets
+from datetime import timedelta
+
+from facefusion.session_manager import clear_session, create_session, get_session, set_session, validate_session
+
+
+def test_get_and_set_session() -> None:
+	session = create_session()
+	access_token = secrets.token_urlsafe(128)
+
+	set_session(access_token, session)
+
+	assert get_session(access_token) == session
+
+
+def test_validate_session() -> None:
+	session = create_session()
+	access_token = secrets.token_urlsafe(128)
+
+	set_session(access_token, session)
+
+	assert validate_session(access_token) is True
+
+	set_session(access_token,
+	{
+		'access_token': session.get('access_token'),
+		'refresh_token': session.get('refresh_token'),
+		'created_at': session.get('created_at'),
+		'expires_at': session.get('expires_at') - timedelta(hours = 1)
+	})
+
+	assert validate_session(access_token) is False
+
+
+def test_clear_session() -> None:
+	session = create_session()
+	access_token = secrets.token_urlsafe(128)
+
+	set_session(access_token, session)
+
+	assert validate_session(access_token) is True
+
+	clear_session(access_token)
+
+	assert validate_session(access_token) is None