⚡️ Speed up function check_api_key by 349%

[codeflash-ai]

📄 349% (3.49x) speedup for check_api_key in cognee/modules/cloud/operations/check_api_key.py

⏱️ Runtime : 150 milliseconds → 33.4 milliseconds (best of 14 runs)

📝 Explanation and details

The optimization introduces a singleton SSL context pattern that eliminates expensive SSL context creation on every function call.

Key Optimization:

• Module-level SSL Context Caching: Instead of calling ssl.create_default_context() on every request, the optimized code creates a single _SECURE_SSL_CONTEXT at module import time and reuses it across all function calls.

Why This Works:
The line profiler reveals that ssl.create_default_context() was consuming 23.1% of total execution time (152ms out of 662ms) in the original code. This function performs expensive operations like loading system certificates, validating trust stores, and configuring SSL protocols. Since SSL context configuration is static and doesn't change between requests, recreating it repeatedly is pure waste.

Performance Impact:

• Runtime improvement: 349% speedup (150ms → 33.4ms)
• Throughput improvement: 27.3% increase (11,165 → 14,210 operations/second)
• The optimized create_secure_ssl_context() function now takes only 0.2ms vs 150ms in the original

Test Case Benefits:
The optimization particularly shines in concurrent scenarios - all throughput tests (small, medium, high volume) benefit significantly because each concurrent call no longer recreates the SSL context. The 200-request high volume test shows substantial gains as the SSL context creation overhead is eliminated across all parallel operations.

Safety: The SSL context is immutable after creation and thread-safe, making this singleton pattern safe for concurrent use in async environments.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 769 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

function to test

import ssl
from unittest.mock import AsyncMock, MagicMock, patch

import aiohttp
import pytest # used for our unit tests
from cognee.modules.cloud.operations.check_api_key import check_api_key

class CloudConnectionError(Exception):
pass
from cognee.modules.cloud.operations.check_api_key import check_api_key

----------------- UNIT TESTS -----------------

Helper to mock aiohttp.ClientSession and its context managers

class MockResponse:
def init(self, status, text=""):
self.status = status
self._text = text

async def text(self):
    return self._text

async def __aenter__(self):
    return self

async def __aexit__(self, exc_type, exc, tb):
    pass

class MockSession:
def init(self, response: MockResponse):
self._response = response

async def __aenter__(self):
    return self

async def __aexit__(self, exc_type, exc, tb):
    pass

def post(self, url, headers):
    # Simulate async context manager for post
    return self._response

---------------- BASIC TEST CASES ----------------

@pytest.mark.asyncio
async def test_check_api_key_success():
"""Test that check_api_key returns None for a successful API key check (HTTP 200)."""
mock_response = MockResponse(status=200)
with patch("aiohttp.ClientSession", return_value=MockSession(mock_response)):
result = await check_api_key("valid_token")

@pytest.mark.asyncio

async def test_check_api_key_many_concurrent_success():
"""Test many concurrent successful API key checks."""
mock_response = MockResponse(status=200)
with patch("aiohttp.ClientSession", return_value=MockSession(mock_response)):
tasks = [check_api_key(f"valid_token_{i}") for i in range(50)]
results = await asyncio.gather(*tasks)

@pytest.mark.asyncio

async def test_check_api_key_throughput_small_load():
"""Throughput test: small load of concurrent requests."""
mock_response = MockResponse(status=200)
with patch("aiohttp.ClientSession", return_value=MockSession(mock_response)):
tasks = [check_api_key(f"token_{i}") for i in range(10)]
results = await asyncio.gather(*tasks)

@pytest.mark.asyncio
async def test_check_api_key_throughput_medium_load():
"""Throughput test: medium load of concurrent requests."""
mock_response = MockResponse(status=200)
with patch("aiohttp.ClientSession", return_value=MockSession(mock_response)):
tasks = [check_api_key(f"token_{i}") for i in range(100)]
results = await asyncio.gather(*tasks)

@pytest.mark.asyncio
async def test_check_api_key_throughput_mixed_load():
"""Throughput test: mix of successful and failed requests."""
# Alternate between success and failure
def session_factory(*args, **kwargs):
class Session:
async def aenter(self): return self
async def aexit(self, exc_type, exc, tb): pass
def post(self, url, headers):
idx = int(headers["X-Api-Key"].split("")[-1])
if idx % 2 == 0:
return MockResponse(status=200)
else:
return MockResponse(status=401, text="Unauthorized")
return Session()
with patch("aiohttp.ClientSession", side_effect=session_factory):
tasks = [check_api_key(f"token{i}") for i in range(50)]
results = await asyncio.gather(*tasks, return_exceptions=True)
# Even indices should succeed, odd should fail
for i, r in enumerate(results):
if i % 2 == 0:
pass
else:
pass

@pytest.mark.asyncio
async def test_check_api_key_throughput_high_volume():
"""Throughput test: high volume of concurrent requests (up to 200)."""
mock_response = MockResponse(status=200)
with patch("aiohttp.ClientSession", return_value=MockSession(mock_response)):
tasks = [check_api_key(f"token_{i}") for i in range(200)]
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.

#------------------------------------------------
import asyncio # used to run async functions

function to test

import ssl
from unittest.mock import AsyncMock, MagicMock, patch

import aiohttp
import pytest # used for our unit tests
from cognee.modules.cloud.operations.check_api_key import check_api_key

class CloudConnectionError(Exception):
pass
from cognee.modules.cloud.operations.check_api_key import check_api_key

-------------------- UNIT TESTS BELOW --------------------

Helper to patch aiohttp.ClientSession for all tests

def mock_response(status=200, text="OK"):
"""
Helper to create a mock aiohttp response object.
"""
response = AsyncMock()
response.status = status
response.text = AsyncMock(return_value=text)
return response

def mock_session_post(status=200, text="OK"):
"""
Helper to patch session.post to return a mock response.
"""
response = mock_response(status, text)
post_cm = AsyncMock()
post_cm.aenter.return_value = response
post_cm.aexit.return_value = None
return post_cm

Basic Test Cases

@pytest.mark.asyncio
async def test_check_api_key_success():
"""
Test that check_api_key returns None for status 200 (success).
"""
with patch("aiohttp.ClientSession") as mock_client_session:
# Mock the session context manager
mock_session = MagicMock()
mock_client_session.return_value.aenter.return_value = mock_session
# Mock the post context manager to return status 200
mock_session.post.return_value = mock_session_post(status=200, text="OK")

    result = await check_api_key("valid_token")

@pytest.mark.asyncio

async def test_check_api_key_many_concurrent_success():
"""
Test many concurrent successful calls.
"""
with patch("aiohttp.ClientSession") as mock_client_session:
mock_session = MagicMock()
mock_client_session.return_value.aenter.return_value = mock_session
mock_session.post.return_value = mock_session_post(status=200, text="OK")

    async def call(token):
        return await check_api_key(token)

    tokens = [f"token_{i}" for i in range(50)]
    results = await asyncio.gather(*(call(token) for token in tokens))

@pytest.mark.asyncio

async def test_check_api_key_throughput_small_load():
"""
Throughput test: small load of concurrent calls.
"""
with patch("aiohttp.ClientSession") as mock_client_session:
mock_session = MagicMock()
mock_client_session.return_value.aenter.return_value = mock_session
mock_session.post.return_value = mock_session_post(status=200, text="OK")

    async def call(token):
        return await check_api_key(token)

    tokens = [f"token_{i}" for i in range(5)]
    results = await asyncio.gather(*(call(token) for token in tokens))

@pytest.mark.asyncio
async def test_check_api_key_throughput_medium_load():
"""
Throughput test: medium load of concurrent calls.
"""
with patch("aiohttp.ClientSession") as mock_client_session:
mock_session = MagicMock()
mock_client_session.return_value.aenter.return_value = mock_session
mock_session.post.return_value = mock_session_post(status=200, text="OK")

    async def call(token):
        return await check_api_key(token)

    tokens = [f"token_{i}" for i in range(50)]
    results = await asyncio.gather(*(call(token) for token in tokens))

@pytest.mark.asyncio
async def test_check_api_key_throughput_high_volume():
"""
Throughput test: high volume concurrent calls (but under 1000).
"""
with patch("aiohttp.ClientSession") as mock_client_session:
mock_session = MagicMock()
mock_client_session.return_value.aenter.return_value = mock_session
mock_session.post.return_value = mock_session_post(status=200, text="OK")

    async def call(token):
        return await check_api_key(token)

    tokens = [f"token_{i}" for i in range(200)]
    results = await asyncio.gather(*(call(token) for token in tokens))

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-check_api_key-mhv3a8g8 and push.