⚡️ Speed up method AstraDBVectorStoreComponent.reset_database_list by 49% in PR #6048 (bugfix-dev-astradb)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #6048

If you approve this dependent PR, these changes will be merged into the original PR branch bugfix-dev-astradb.

This PR will be automatically closed if the original PR is merged.

---

📄 49% (0.49x) speedup for AstraDBVectorStoreComponent.reset_database_list in src/backend/base/langflow/components/vectorstores/astradb.py

⏱️ Runtime : 974 microseconds → 655 microseconds (best of 113 runs)

📝 Explanation and details

To optimize the provided Python code for better performance, we can adopt the following strategies.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

import pytest  # used for our unit tests
# function to test
from langflow.base.vectorstores.model import LCVectorStoreComponent
from langflow.components.vectorstores.astradb import \
    AstraDBVectorStoreComponent

# unit tests

# Mock class to simulate get_database_list method
class MockAstraDBVectorStoreComponent(AstraDBVectorStoreComponent):
    def __init__(self, database_list):
        self._database_list = database_list

    def get_database_list(self):
        return self._database_list


def test_basic_single_database_option():
    # Test with a single database option
    component = MockAstraDBVectorStoreComponent({
        "db1": {"collections": ["col1"], "api_endpoint": "http://example.com"}
    })
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    codeflash_output = component.reset_database_list(build_config)


def test_basic_multiple_database_options():
    # Test with multiple database options
    component = MockAstraDBVectorStoreComponent({
        "db1": {"collections": ["col1"], "api_endpoint": "http://example.com"},
        "db2": {"collections": ["col2"], "api_endpoint": "http://example2.com"}
    })
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    codeflash_output = component.reset_database_list(build_config)


def test_edge_empty_database_list():
    # Test with an empty database list
    component = MockAstraDBVectorStoreComponent({})
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    codeflash_output = component.reset_database_list(build_config)


def test_edge_missing_fields():
    # Test with database list missing fields
    component = MockAstraDBVectorStoreComponent({
        "db1": {"collections": ["col1"]}  # Missing "api_endpoint"
    })
    with pytest.raises(ValueError):
        build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
        component.reset_database_list(build_config)


def test_edge_extra_fields():
    # Test with database list with extra fields
    component = MockAstraDBVectorStoreComponent({
        "db1": {"collections": ["col1"], "api_endpoint": "http://example.com", "extra_field": "extra_value"}
    })
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    codeflash_output = component.reset_database_list(build_config)


def test_error_handling_exception_during_initialization():
    # Test with an exception during database initialization
    class ExceptionMockAstraDBVectorStoreComponent(AstraDBVectorStoreComponent):
        def get_database_list(self):
            raise Exception("Test Exception")

    component = ExceptionMockAstraDBVectorStoreComponent()
    with pytest.raises(ValueError) as excinfo:
        build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
        component.reset_database_list(build_config)


def test_large_scale_many_databases():
    # Test with a large number of databases
    database_list = {f"db{i}": {"collections": [f"col{i}"], "api_endpoint": f"http://example{i}.com"} for i in range(1000)}
    component = MockAstraDBVectorStoreComponent(database_list)
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    codeflash_output = component.reset_database_list(build_config)


def test_special_characters_in_database_names():
    # Test with special characters in database names
    component = MockAstraDBVectorStoreComponent({
        "db!@#": {"collections": ["col1"], "api_endpoint": "http://example.com"}
    })
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    codeflash_output = component.reset_database_list(build_config)


def test_nested_structures_in_database_list():
    # Test with nested structures in database list
    component = MockAstraDBVectorStoreComponent({
        "db1": {"collections": [{"name": "col1"}], "api_endpoint": "http://example.com"}
    })
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    codeflash_output = component.reset_database_list(build_config)




import pytest  # used for our unit tests
# function to test
from langflow.base.vectorstores.model import LCVectorStoreComponent
from langflow.components.vectorstores.astradb import \
    AstraDBVectorStoreComponent


# unit tests
class MockAstraDBVectorStoreComponent(AstraDBVectorStoreComponent):
    def __init__(self, mock_data):
        self.mock_data = mock_data

    def get_database_list(self):
        return self.mock_data

def test_basic_functionality():
    # Standard Input
    mock_data = {
        "db1": {"collections": ["col1"], "api_endpoint": "endpoint1"},
        "db2": {"collections": ["col2"], "api_endpoint": "endpoint2"}
    }
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    expected_config = {
        "api_endpoint": {
            "options": ["db1", "db2"],
            "options_metadata": [
                {"collections": ["col1"], "api_endpoint": "endpoint1"},
                {"collections": ["col2"], "api_endpoint": "endpoint2"}
            ],
            "value": ""
        }
    }
    codeflash_output = component.reset_database_list(build_config)

def test_empty_database_list():
    # Empty Database List
    mock_data = {}
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    expected_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    codeflash_output = component.reset_database_list(build_config)

def test_no_api_endpoint_key():
    # No `api_endpoint` Key in `build_config`
    mock_data = {}
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {}
    with pytest.raises(KeyError):
        component.reset_database_list(build_config)

def test_no_options_key():
    # No `options` Key in `api_endpoint`
    mock_data = {
        "db1": {"collections": ["col1"], "api_endpoint": "endpoint1"}
    }
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {"api_endpoint": {"options_metadata": [], "value": ""}}
    expected_config = {
        "api_endpoint": {
            "options": ["db1"],
            "options_metadata": [{"collections": ["col1"], "api_endpoint": "endpoint1"}],
            "value": ""
        }
    }
    codeflash_output = component.reset_database_list(build_config)

def test_no_options_metadata_key():
    # No `options_metadata` Key in `api_endpoint`
    mock_data = {
        "db1": {"collections": ["col1"], "api_endpoint": "endpoint1"}
    }
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {"api_endpoint": {"options": [], "value": ""}}
    expected_config = {
        "api_endpoint": {
            "options": ["db1"],
            "options_metadata": [{"collections": ["col1"], "api_endpoint": "endpoint1"}],
            "value": ""
        }
    }
    codeflash_output = component.reset_database_list(build_config)

def test_exception_in_initialize_database_options():
    # Exception in `_initialize_database_options`
    def mock_get_database_list():
        raise Exception("Mock exception")
    component = MockAstraDBVectorStoreComponent({})
    component.get_database_list = mock_get_database_list
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    with pytest.raises(ValueError, match="Error fetching database options: Mock exception"):
        component.reset_database_list(build_config)

def test_large_number_of_databases():
    # Large Number of Databases
    mock_data = {f"db{i}": {"collections": [f"col{i}"], "api_endpoint": f"endpoint{i}"} for i in range(1000)}
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    expected_config = {
        "api_endpoint": {
            "options": [f"db{i}" for i in range(1000)],
            "options_metadata": [
                {"collections": [f"col{i}"], "api_endpoint": f"endpoint{i}"} for i in range(1000)
            ],
            "value": ""
        }
    }
    codeflash_output = component.reset_database_list(build_config)

def test_special_characters_in_database_names():
    # Special Characters in Database Names
    mock_data = {
        "db@1!": {"collections": ["col1"], "api_endpoint": "endpoint1"}
    }
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    expected_config = {
        "api_endpoint": {
            "options": ["db@1!"],
            "options_metadata": [{"collections": ["col1"], "api_endpoint": "endpoint1"}],
            "value": ""
        }
    }
    codeflash_output = component.reset_database_list(build_config)

def test_unicode_characters_in_database_names():
    # Unicode Characters
    mock_data = {
        "数据库": {"collections": ["集合"], "api_endpoint": "端点"}
    }
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    expected_config = {
        "api_endpoint": {
            "options": ["数据库"],
            "options_metadata": [{"collections": ["集合"], "api_endpoint": "端点"}],
            "value": ""
        }
    }
    codeflash_output = component.reset_database_list(build_config)

def test_nested_collections():
    # Nested Collections
    mock_data = {
        "db1": {"collections": [["col1", "col2"]], "api_endpoint": "endpoint1"}
    }
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    expected_config = {
        "api_endpoint": {
            "options": ["db1"],
            "options_metadata": [{"collections": [["col1", "col2"]], "api_endpoint": "endpoint1"}],
            "value": ""
        }
    }
    codeflash_output = component.reset_database_list(build_config)

def test_mixed_data_types_in_collections():
    # Mixed Data Types in `collections`
    mock_data = {
        "db1": {"collections": ["col1", 2, ["col3"]], "api_endpoint": "endpoint1"}
    }
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    expected_config = {
        "api_endpoint": {
            "options": ["db1"],
            "options_metadata": [{"collections": ["col1", 2, ["col3"]], "api_endpoint": "endpoint1"}],
            "value": ""
        }
    }
    codeflash_output = component.reset_database_list(build_config)




def test_verify_side_effects():
    # Verify Side Effects
    mock_data = {
        "db1": {"collections": ["col1"], "api_endpoint": "endpoint1"}
    }
    component = MockAstraDBVectorStoreComponent(mock_data)
    build_config = {"api_endpoint": {"options": [], "options_metadata": [], "value": ""}}
    component.reset_database_list(build_config)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.