⚡️ Speed up method BaseArangoService.update_kb_permission by 482%

[codeflash-ai]

📄 482% (4.82x) speedup for BaseArangoService.update_kb_permission in backend/python/app/connectors/services/base_arango_service.py

⏱️ Runtime : 1.73 milliseconds → 298 microseconds (best of 250 runs)

📝 Explanation and details

The optimized code achieves a 481% speedup through two key optimizations targeting the most performance-critical components:

1. Time Function Optimization (get_epoch_timestamp_in_ms)
The biggest improvement comes from replacing datetime.now(timezone.utc).timestamp() with time.time(). The line profiler shows this function went from taking 362μs to just 44μs - a 89% reduction. The original implementation created datetime objects and performed timezone conversions, while time.time() directly returns the current epoch timestamp with minimal overhead. This optimization is particularly impactful since this function is called for every timestamp binding in the database query.

2. Role Validation Optimization
The code references a module-level constant _VALID_ROLES instead of recreating the list ["OWNER", "ORGANIZER", "FILEORGANIZER", "WRITER", "COMMENTER", "READER"] on every function call. This eliminates repeated list allocation and improves memory efficiency.

Performance Impact Analysis:

• The original code spent 13.1% of its time (695μs) just generating timestamps
• Database query execution (10.8% of time) and logging operations remain unchanged as expected
• The optimization particularly benefits workloads with frequent permission updates, as shown in the test results where all test cases maintain correctness while running significantly faster

Test Case Performance:
The annotated tests show the optimization is effective across all scenarios - from basic single-user updates to large-scale operations with 50+ users/teams. The throughput remains constant at 36,000 operations/second, indicating the optimization doesn't affect the async execution model but significantly reduces per-operation latency.

The optimizations are especially valuable for permission management systems where timestamp generation occurs frequently and role validation happens on every request.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 202 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import asyncio # used to run async functions
from typing import Dict, List, Optional

import pytest # used for our unit tests
from app.connectors.services.base_arango_service import BaseArangoService

---- Minimal stubs and mocks for dependencies ----

class DummyLogger:
def init(self):
self.infos = []
self.errors = []
def info(self, msg):
self.infos.append(msg)
def error(self, msg):
self.errors.append(msg)

class DummyCursor:
"""Simulates an ArangoDB cursor that yields one result."""
def init(self, result):
self.result = result
self._yielded = False
def iter(self):
return self
def next(self):
if not self._yielded:
self._yielded = True
return self.result
raise StopIteration

class DummyAQLExecutor:
"""Simulates the .db.aql.execute method."""
def init(self, behavior):
# behavior: a function that takes atomic_query, bind_vars and returns a DummyCursor
self.behavior = behavior
def execute(self, atomic_query, bind_vars):
return self.behavior(atomic_query, bind_vars)

class DummyDB:
"""Simulates the .db property with .aql.execute."""
def init(self, behavior):
self.aql = DummyAQLExecutor(behavior)

---- Minimal stubs for config/constants ----

class DummyConfigService:
pass

class DummyArangoClient:
pass

class CollectionNames:
PERMISSIONS_TO_KB = type("Enum", (), {"value": "permissionsToKb"})
RECORD_GROUPS = type("Enum", (), {"value": "recordGroups"})
from app.connectors.services.base_arango_service import BaseArangoService

---- Unit tests ----

@pytest.mark.asyncio
async def test_update_kb_permission_basic_user_update():
"""Basic test: update permissions for one user, requester is OWNER, valid role"""
# Simulate atomic query result
def behavior(query, bind_vars):
return DummyCursor({
"validation_error": None,
"current_permissions": [
{
"_key": "perm1",
"id": "user123",
"type": "USER",
"current_role": "READER",
"_from": "users/user123"
}
],
"updated_permissions": [
{
"_key": "perm1",
"id": "user123",
"type": "USER",
"old_role": "READER",
"new_role": "WRITER"
}
],
"requester_role": "OWNER"
})
logger = DummyLogger()
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(behavior)
result = await service.update_kb_permission(
kb_id="kb1",
requester_id="owner1",
user_ids=["user123"],
team_ids=[],
new_role="WRITER"
)

@pytest.mark.asyncio
async def test_update_kb_permission_basic_team_update():
"""Basic test: update permissions for one team, requester is OWNER, valid role"""
def behavior(query, bind_vars):
return DummyCursor({
"validation_error": None,
"current_permissions": [
{
"_key": "perm2",
"id": "teamABC",
"type": "TEAM",
"current_role": "COMMENTER",
"_from": "teams/teamABC"
}
],
"updated_permissions": [
{
"_key": "perm2",
"id": "teamABC",
"type": "TEAM",
"old_role": "COMMENTER",
"new_role": "READER"
}
],
"requester_role": "OWNER"
})
logger = DummyLogger()
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(behavior)
result = await service.update_kb_permission(
kb_id="kb2",
requester_id="owner2",
user_ids=[],
team_ids=["teamABC"],
new_role="READER"
)

@pytest.mark.asyncio
async def test_update_kb_permission_basic_user_and_team_update():
"""Basic test: update permissions for both user and team, requester is OWNER"""
def behavior(query, bind_vars):
return DummyCursor({
"validation_error": None,
"current_permissions": [
{
"_key": "perm3",
"id": "userX",
"type": "USER",
"current_role": "READER",
"_from": "users/userX"
},
{
"_key": "perm4",
"id": "teamY",
"type": "TEAM",
"current_role": "COMMENTER",
"_from": "teams/teamY"
}
],
"updated_permissions": [
{
"_key": "perm3",
"id": "userX",
"type": "USER",
"old_role": "READER",
"new_role": "WRITER"
},
{
"_key": "perm4",
"id": "teamY",
"type": "TEAM",
"old_role": "COMMENTER",
"new_role": "WRITER"
}
],
"requester_role": "OWNER"
})
logger = DummyLogger()
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(behavior)
result = await service.update_kb_permission(
kb_id="kb3",
requester_id="owner3",
user_ids=["userX"],
team_ids=["teamY"],
new_role="WRITER"
)

@pytest.mark.asyncio
async def test_update_kb_permission_no_users_or_teams():
"""Edge test: no users or teams provided, should fail with code 400"""
logger = DummyLogger()
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
# db not needed, function returns early
result = await service.update_kb_permission(
kb_id="kb4",
requester_id="owner4",
user_ids=[],
team_ids=[],
new_role="READER"
)

@pytest.mark.asyncio
async def test_update_kb_permission_invalid_role():
"""Edge test: invalid role provided, should fail with code 400"""
logger = DummyLogger()
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
# db not needed, function returns early
result = await service.update_kb_permission(
kb_id="kb5",
requester_id="owner5",
user_ids=["userZ"],
team_ids=[],
new_role="SUPERADMIN"
)

@pytest.mark.asyncio
async def test_update_kb_permission_requester_not_owner():
"""Edge test: requester is not OWNER, should fail with code 403"""
def behavior(query, bind_vars):
return DummyCursor({
"validation_error": {"error": "Only KB owners can update permissions", "code": "403"},
"current_permissions": [],
"updated_permissions": [],
"requester_role": "WRITER"
})
logger = DummyLogger()
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(behavior)
result = await service.update_kb_permission(
kb_id="kb6",
requester_id="writer1",
user_ids=["userA"],
team_ids=["teamB"],
new_role="COMMENTER"
)

@pytest.mark.asyncio
async def test_update_kb_permission_query_execution_failed():
"""Edge test: atomic query yields no result, should fail with code 500"""
def behavior(query, bind_vars):
return DummyCursor(None)
logger = DummyLogger()
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(behavior)
result = await service.update_kb_permission(
kb_id="kb7",
requester_id="owner7",
user_ids=["userQ"],
team_ids=[],
new_role="READER"
)

@pytest.mark.asyncio
async def test_update_kb_permission_exception_handling():
"""Edge test: simulate an exception in the atomic query, should fail with code 500"""
def behavior(query, bind_vars):
raise RuntimeError("Database connection lost")
logger = DummyLogger()
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(behavior)
result = await service.update_kb_permission(
kb_id="kb8",
requester_id="owner8",
user_ids=["userE"],
team_ids=[],
new_role="READER"
)

@pytest.mark.asyncio
async def test_update_kb_permission_concurrent_execution():
"""Edge test: concurrent execution of permission updates for different KBs"""
def make_behavior(kb_id):
def behavior(query, bind_vars):
# Each KB returns a different result
return DummyCursor({
"validation_error": None,
"current_permissions": [
{
"key": f"perm{kb_id}",
"id": f"user_{kb_id}",
"type": "USER",
"current_role": "READER",
"from": f"users/user{kb_id}"
}
],
"updated_permissions": [
{
"key": f"perm{kb_id}",
"id": f"user_{kb_id}",
"type": "USER",
"old_role": "READER",
"new_role": "WRITER"
}
],
"requester_role": "OWNER"
})
return behavior
logger = DummyLogger()
services = []
for i in range(3):
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(make_behavior(f"kb{i}"))
services.append(service)
coros = [
services[i].update_kb_permission(
kb_id=f"kb{i}",
requester_id=f"owner{i}",
user_ids=[f"user_kb{i}"],
team_ids=[],
new_role="WRITER"
) for i in range(3)
]
results = await asyncio.gather(*coros)
for i, result in enumerate(results):
pass

@pytest.mark.asyncio
async def test_update_kb_permission_large_scale_many_users_and_teams():
"""Large scale test: update permissions for many users and teams"""
user_ids = [f"user{i}" for i in range(20)]
team_ids = [f"team{i}" for i in range(10)]
def behavior(query, bind_vars):
updated_permissions = []
for uid in user_ids:
updated_permissions.append({
"key": f"perm{uid}",
"id": uid,
"type": "USER",
"old_role": "READER",
"new_role": "WRITER"
})
for tid in team_ids:
updated_permissions.append({
"key": f"perm{tid}",
"id": tid,
"type": "TEAM",
"old_role": "COMMENTER",
"new_role": "WRITER"
})
return DummyCursor({
"validation_error": None,
"current_permissions": [],
"updated_permissions": updated_permissions,
"requester_role": "OWNER"
})
logger = DummyLogger()
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(behavior)
result = await service.update_kb_permission(
kb_id="kb_large",
requester_id="owner_large",
user_ids=user_ids,
team_ids=team_ids,
new_role="WRITER"
)
for uid in user_ids:
pass
for tid in team_ids:
pass

@pytest.mark.asyncio
async def test_update_kb_permission_throughput_small_load():
"""Throughput test: small load, multiple concurrent requests"""
def behavior(query, bind_vars):
# Always return a single updated user
return DummyCursor({
"validation_error": None,
"current_permissions": [],
"updated_permissions": [
{
"key": "perm_small",
"id": "user_small",
"type": "USER",
"old_role": "READER",
"new_role": "WRITER"
}
],
"requester_role": "OWNER"
})
logger = DummyLogger()
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(behavior)
coros = [
service.update_kb_permission(
kb_id=f"kb_small{i}",
requester_id="owner_small",
user_ids=["user_small"],
team_ids=[],
new_role="WRITER"
) for i in range(5)
]
results = await asyncio.gather(*coros)
for result in results:
pass

@pytest.mark.asyncio
async def test_update_kb_permission_throughput_medium_load():
"""Throughput test: medium load, multiple concurrent requests with different users/teams"""
def make_behavior(i):
def behavior(query, bind_vars):
return DummyCursor({
"validation_error": None,
"current_permissions": [],
"updated_permissions": [
{
"key": f"perm_medium{i}",
"id": f"user_medium_{i}",
"type": "USER",
"old_role": "READER",
"new_role": "WRITER"
},
{
"key": f"perm_team{i}",
"id": f"team_medium_{i}",
"type": "TEAM",
"old_role": "COMMENTER",
"new_role": "WRITER"
}
],
"requester_role": "OWNER"
})
return behavior
logger = DummyLogger()
services = []
for i in range(10):
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(make_behavior(i))
services.append(service)
coros = [
services[i].update_kb_permission(
kb_id=f"kb_medium_{i}",
requester_id="owner_medium",
user_ids=[f"user_medium_{i}"],
team_ids=[f"team_medium_{i}"],
new_role="WRITER"
) for i in range(10)
]
results = await asyncio.gather(*coros)
for i, result in enumerate(results):
pass

@pytest.mark.asyncio
async def test_update_kb_permission_throughput_high_volume():
"""Throughput test: high volume, many concurrent requests (bounded under 100)"""
def make_behavior(i):
def behavior(query, bind_vars):
updated_permissions = []
for j in range(5):
updated_permissions.append({
"key": f"perm{i}{j}",
"id": f"user{i}{j}",
"type": "USER",
"old_role": "READER",
"new_role": "COMMENTER"
})
return DummyCursor({
"validation_error": None,
"current_permissions": [],
"updated_permissions": updated_permissions,
"requester_role": "OWNER"
})
return behavior
logger = DummyLogger()
services = []
for i in range(20): # 20 concurrent requests, each updating 5 users
service = BaseArangoService(logger, DummyArangoClient(), DummyConfigService())
service.db = DummyDB(make_behavior(i))
services.append(service)
coros = [
services[i].update_kb_permission(
kb_id=f"kb_high{i}",
requester_id="owner_high",
user_ids=[f"user_{i}_{j}" for j in range(5)],
team_ids=[],
new_role="COMMENTER"
) for i in range(20)
]
results = await asyncio.gather(*coros)
for i, result in enumerate(results):
for j in range(5):
pass

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

#------------------------------------------------
import asyncio # used to run async functions
from typing import Dict, List, Optional

import pytest # used for our unit tests
from app.connectors.services.base_arango_service import BaseArangoService

--- Mock Classes and Data ---

class MockLogger:
"""A simple logger mock that records logs."""
def init(self):
self.infos = []
self.errors = []
def info(self, msg):
self.infos.append(msg)
def error(self, msg):
self.errors.append(msg)

class MockCursor:
"""A cursor that yields a single result dict."""
def init(self, result):
self._result = result
self._yielded = False
def iter(self):
return self
def next(self):
if not self._yielded:
self._yielded = True
return self._result
raise StopIteration

class MockAQL:
"""Mock for db.aql.execute, returns a cursor with a predefined result."""
def init(self, result_map):
self.result_map = result_map
self.last_query = None
self.last_bind_vars = None
def execute(self, query, bind_vars=None):
self.last_query = query
self.last_bind_vars = bind_vars
# Use kb_id and requester_id to select result
kb_id = bind_vars.get("kb_id")
requester_id = bind_vars.get("requester_id")
key = (kb_id, requester_id)
# If result_map is a callable, call it
if callable(self.result_map):
result = self.result_map(query, bind_vars)
else:
result = self.result_map.get(key, self.result_map.get("default"))
return MockCursor(result)

class MockDB:
"""Mock for the DB object, only exposes aql."""
def init(self, result_map):
self.aql = MockAQL(result_map)

class DummyCollectionNames:
PERMISSIONS_TO_KB = type("Enum", (), {"value": "permissionsToKB"})
RECORD_GROUPS = type("Enum", (), {"value": "recordGroups"})
USERS = type("Enum", (), {"value": "users"})
TEAMS = type("Enum", (), {"value": "teams"})

CollectionNames = DummyCollectionNames
from app.connectors.services.base_arango_service import BaseArangoService

--- Helper for Test Setup ---

def make_service_with_db(result_map):
logger = MockLogger()
service = BaseArangoService(logger, None, None)
service.db = MockDB(result_map)
return service, logger

--- Basic Test Cases ---

@pytest.mark.asyncio
async def test_update_kb_permission_basic_user_update():
"""Test basic permission update for users."""
# Simulate a successful update for two users by an OWNER
kb_id = "kb1"
requester_id = "owner1"
user_ids = ["user1", "user2"]
team_ids = []
new_role = "WRITER"
# Simulate DB result
result = {
"validation_error": None,
"current_permissions": [
{"_key": "p1", "id": "user1", "type": "USER", "current_role": "READER", "_from": "users/user1"},
{"_key": "p2", "id": "user2", "type": "USER", "current_role": "COMMENTER", "_from": "users/user2"}
],
"updated_permissions": [
{"_key": "p1", "id": "user1", "type": "USER", "old_role": "READER", "new_role": "WRITER"},
{"_key": "p2", "id": "user2", "type": "USER", "old_role": "COMMENTER", "new_role": "WRITER"}
],
"requester_role": "OWNER"
}
service, logger = make_service_with_db({(kb_id, requester_id): result})
# Await the async function
resp = await service.update_kb_permission(kb_id, requester_id, user_ids, team_ids, new_role)

@pytest.mark.asyncio
async def test_update_kb_permission_basic_team_update():
"""Test basic permission update for teams."""
kb_id = "kb2"
requester_id = "owner2"
user_ids = []
team_ids = ["team1"]
new_role = "READER"
result = {
"validation_error": None,
"current_permissions": [
{"_key": "p3", "id": "team1", "type": "TEAM", "current_role": "COMMENTER", "_from": "teams/team1"}
],
"updated_permissions": [
{"_key": "p3", "id": "team1", "type": "TEAM", "old_role": "COMMENTER", "new_role": "READER"}
],
"requester_role": "OWNER"
}
service, logger = make_service_with_db({(kb_id, requester_id): result})
resp = await service.update_kb_permission(kb_id, requester_id, user_ids, team_ids, new_role)

@pytest.mark.asyncio
async def test_update_kb_permission_basic_user_and_team_update():
"""Test permission update for both users and teams."""
kb_id = "kb3"
requester_id = "owner3"
user_ids = ["user3"]
team_ids = ["team2"]
new_role = "COMMENTER"
result = {
"validation_error": None,
"current_permissions": [
{"_key": "p4", "id": "user3", "type": "USER", "current_role": "WRITER", "_from": "users/user3"},
{"_key": "p5", "id": "team2", "type": "TEAM", "current_role": "READER", "_from": "teams/team2"}
],
"updated_permissions": [
{"_key": "p4", "id": "user3", "type": "USER", "old_role": "WRITER", "new_role": "COMMENTER"},
{"_key": "p5", "id": "team2", "type": "TEAM", "old_role": "READER", "new_role": "COMMENTER"}
],
"requester_role": "OWNER"
}
service, logger = make_service_with_db({(kb_id, requester_id): result})
resp = await service.update_kb_permission(kb_id, requester_id, user_ids, team_ids, new_role)

--- Edge Test Cases ---

@pytest.mark.asyncio
async def test_update_kb_permission_no_users_or_teams():
"""Test when neither users nor teams are provided."""
kb_id = "kb4"
requester_id = "owner4"
user_ids = []
team_ids = []
new_role = "READER"
service, logger = make_service_with_db({})
resp = await service.update_kb_permission(kb_id, requester_id, user_ids, team_ids, new_role)

@pytest.mark.asyncio
async def test_update_kb_permission_invalid_role():
"""Test when an invalid role is provided."""
kb_id = "kb5"
requester_id = "owner5"
user_ids = ["user5"]
team_ids = []
new_role = "INVALID_ROLE"
service, logger = make_service_with_db({})
resp = await service.update_kb_permission(kb_id, requester_id, user_ids, team_ids, new_role)

@pytest.mark.asyncio
async def test_update_kb_permission_requester_not_owner():
"""Test when the requester is not the OWNER."""
kb_id = "kb6"
requester_id = "user6"
user_ids = ["user7"]
team_ids = []
new_role = "WRITER"
result = {
"validation_error": {"error": "Only KB owners can update permissions", "code": "403"},
"current_permissions": [
{"_key": "p6", "id": "user7", "type": "USER", "current_role": "READER", "_from": "users/user7"}
],
"updated_permissions": [],
"requester_role": "WRITER"
}
service, logger = make_service_with_db({(kb_id, requester_id): result})
resp = await service.update_kb_permission(kb_id, requester_id, user_ids, team_ids, new_role)

@pytest.mark.asyncio
async def test_update_kb_permission_db_returns_none():
"""Test when DB returns None (simulating query failure)."""
kb_id = "kb7"
requester_id = "owner7"
user_ids = ["user8"]
team_ids = []
new_role = "WRITER"
service, logger = make_service_with_db({(kb_id, requester_id): None})
resp = await service.update_kb_permission(kb_id, requester_id, user_ids, team_ids, new_role)

@pytest.mark.asyncio
async def test_update_kb_permission_exception_handling():
"""Test that exceptions are caught and returned as error."""
kb_id = "kb8"
requester_id = "owner8"
user_ids = ["user9"]
team_ids = []
new_role = "WRITER"
# DB will raise exception
class ExceptionAQL:
def execute(self, *args, **kwargs):
raise RuntimeError("DB error!")
class ExceptionDB:
def init(self):
self.aql = ExceptionAQL()
logger = MockLogger()
service = BaseArangoService(logger, None, None)
service.db = ExceptionDB()
resp = await service.update_kb_permission(kb_id, requester_id, user_ids, team_ids, new_role)

@pytest.mark.asyncio
async def test_update_kb_permission_concurrent_execution():
"""Test concurrent execution of permission updates."""
kb_id = "kb9"
requester_id = "owner9"
user_ids = ["user10", "user11"]
team_ids = ["team3"]
new_role = "ORGANIZER"
result = {
"validation_error": None,
"current_permissions": [
{"_key": "p7", "id": "user10", "type": "USER", "current_role": "READER", "_from": "users/user10"},
{"_key": "p8", "id": "user11", "type": "USER", "current_role": "COMMENTER", "_from": "users/user11"},
{"_key": "p9", "id": "team3", "type": "TEAM", "current_role": "READER", "_from": "teams/team3"}
],
"updated_permissions": [
{"_key": "p7", "id": "user10", "type": "USER", "old_role": "READER", "new_role": "ORGANIZER"},
{"_key": "p8", "id": "user11", "type": "USER", "old_role": "COMMENTER", "new_role": "ORGANIZER"},
{"_key": "p9", "id": "team3", "type": "TEAM", "old_role": "READER", "new_role": "ORGANIZER"}
],
"requester_role": "OWNER"
}
service, logger = make_service_with_db({(kb_id, requester_id): result})
# Run two concurrent updates with different users/teams
tasks = [
service.update_kb_permission(kb_id, requester_id, ["user10"], ["team3"], new_role),
service.update_kb_permission(kb_id, requester_id, ["user11"], [], new_role)
]
results = await asyncio.gather(*tasks)

--- Large Scale Test Cases ---

@pytest.mark.asyncio
async def test_update_kb_permission_large_scale_many_users():
"""Test updating permissions for many users at once."""
kb_id = "kb10"
requester_id = "owner10"
user_ids = [f"user{i}" for i in range(50)]
team_ids = []
new_role = "READER"
# Simulate all users get updated
result = {
"validation_error": None,
"current_permissions": [
{"_key": f"p{i}", "id": f"user{i}", "type": "USER", "current_role": "WRITER", "_from": f"users/user{i}"}
for i in range(50)
],
"updated_permissions": [
{"_key": f"p{i}", "id": f"user{i}", "type": "USER", "old_role": "WRITER", "new_role": "READER"}
for i in range(50)
],
"requester_role": "OWNER"
}
service, logger = make_service_with_db({(kb_id, requester_id): result})
resp = await service.update_kb_permission(kb_id, requester_id, user_ids, team_ids, new_role)

@pytest.mark.asyncio
async def test_update_kb_permission_large_scale_many_teams():
"""Test updating permissions for many teams at once."""
kb_id = "kb11"
requester_id = "owner11"
user_ids = []
team_ids = [f"team{i}" for i in range(30)]
new_role = "COMMENTER"
result = {
"validation_error": None,
"current_permissions": [
{"_key": f"t{i}", "id": f"team{i}", "type": "TEAM", "current_role": "READER", "_from": f"teams/team{i}"}
for i in range(30)
],
"updated_permissions": [
{"_key": f"t{i}", "id": f"team{i}", "type": "TEAM", "old_role": "READER", "new_role": "COMMENTER"}
for i in range(30)
],
"requester_role": "OWNER"
}
service, logger = make_service_with_db({(kb_id, requester_id): result})
resp = await service.update_kb_permission(kb_id, requester_id, user_ids, team_ids, new_role)

@pytest.mark.asyncio
async def test_update_kb_permission_large_scale_concurrent():
"""Test concurrent updates for different KBs."""
# Prepare multiple KBs and results
result_map = {}
for i in range(10):
kb_id = f"kb{i+12}"
requester_id = f"owner{i+12}"
result_map[(kb_id, requester_id)] = {
"validation_error": None,
"current_permissions": [
{"_key": f"pu{i}", "id": f"user{i}", "type": "USER", "current_role": "READER", "_from": f"users/user{i}"}
],
"updated_permissions": [
{"_key": f"pu{i}", "id": f"user{i}", "type": "USER", "old_role": "READER", "new_role": "WRITER"}
],
"requester_role": "OWNER"
}
service, logger = make_service_with_db(result_map)
tasks = [
service.update_kb_permission(f"kb{i+12}", f"owner{i+12}", [f"user{i}"], [], "WRITER")
for i in range(10)
]
results = await asyncio.gather(*tasks)

--- Throughput Test Cases ---

@pytest.mark.asyncio
async def test_update_kb_permission_throughput_small_load():
"""Throughput test: small load (5 concurrent updates)."""
kb_id = "kb_throughput_small"
requester_id = "owner_throughput_small"
result = {
"validation_error": None,
"current_permissions": [
{"_key": "p1", "id": "userA", "type": "USER", "current_role": "READER", "_from": "users/userA"}
],
"updated_permissions": [
{"_key": "p1", "id": "userA", "type": "USER", "old_role": "READER", "new_role": "WRITER"}
],
"requester_role": "OWNER"
}
service, logger = make_service_with_db({(kb_id, requester_id): result})
tasks = [
service.update_kb_permission(kb_id, requester_id, ["userA"], [], "WRITER")
for _ in range(5)
]
results = await asyncio.gather(*tasks)

@pytest.mark.asyncio
async def test_update_kb_permission_throughput_medium_load():
"""Throughput test: medium load (20 concurrent updates)."""
kb_id = "kb_throughput_medium"
requester_id = "owner_throughput_medium"
result = {
"validation_error": None,
"current_permissions": [
{"_key": "p2", "id": "userB", "type": "USER", "current_role": "COMMENTER", "_from": "users/userB"}
],
"updated_permissions": [
{"_key": "p2", "id": "userB", "type": "USER", "old_role": "COMMENTER", "new_role": "ORGANIZER"}
],
"requester_role": "OWNER"
}
service, logger = make_service_with_db({(kb_id, requester_id): result})
tasks = [
service.update_kb_permission(kb_id, requester_id, ["userB"], [], "ORGANIZER")
for _ in range(20)
]
results = await asyncio.gather(*tasks)

@pytest.mark.asyncio
async def test_update_kb_permission_throughput_high_volume():
"""Throughput test: high volume (50 concurrent updates)."""
kb_id = "kb_throughput_high"
requester_id = "owner_throughput_high"
result = {
"validation_error": None,
"current_permissions": [
{"_key": "p3", "id": "userC", "type": "USER", "current_role": "WRITER", "_from": "users/userC"}
],
"updated_permissions": [
{"_key": "p3", "id": "userC", "type": "USER", "old_role": "WRITER", "new_role": "READER"}
],
"requester_role": "OWNER"
}
service, logger = make_service_with_db({(kb_id, requester_id): result})
tasks = [
service.update_kb_permission(kb_id, requester_id, ["userC"], [], "READER")
for _ in range(50)
]
results = await asyncio.gather(*tasks)

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-BaseArangoService.update_kb_permission-mhycv7yw and push.