feat: add dialog support to dropdown components and enhance AstraDB i…

Open deon-sanchez opened this issue 1 year ago • 1 comments

…ntegration

Introduced a new NodeDialog component for improved user interaction when connecting to Astra DB.
Updated AstraDBVectorStoreComponent to include a dialog option for creating new collections.
Enhanced dropdown components to support dialog functionality, allowing users to create new options directly from the dropdown.
Added hasDialog property to relevant component types to manage dialog visibility.
Refactored utility functions to support new naming conventions and improve code clarity.

Dec 18 '24 23:12 deon-sanchez

CodSpeed Performance Report

Merging #5349 will degrade performances by 91.11%

_{Comparing feat-common-dialog-astra (73d0476) with main (f08c18f)}

Summary

❌ 3 regressions
✅ 6 untouched benchmarks

:warning: Please fix the performance issues or acknowledge them on CodSpeed.

Benchmarks breakdown

	Benchmark	`main`	`feat-common-dialog-astra`	Change
❌	`test_build_flow`	392 ms	4,411.4 ms	-91.11%
❌	`test_invalid_run_with_input_type_chat`	25.6 ms	29.7 ms	-13.71%
❌	`test_successful_run_with_output_type_debug`	181.1 ms	251.4 ms	-27.95%

Dec 18 '24 23:12 codspeed-hq[bot]

⚡️ Codeflash found optimizations for this PR

📄 27% (0.27x) speedup for `AstraDBVectorStoreComponent.get_database_list` in `src/backend/base/langflow/components/vectorstores/astradb.py`

⏱️ Runtime : 19.3 milliseconds → 15.2 milliseconds (best of 5 runs)

📝 Explanation and details

These changes aim to cut down on redundant operations and enhance the efficiency of the program.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from unittest.mock import MagicMock

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

# unit tests

@pytest.fixture
def mock_data_api_client(mocker):
    mock_client = mocker.patch('astrapy.DataAPIClient')
    mock_admin_client = MagicMock()
    mock_client.return_value.get_admin.return_value = mock_admin_client
    return mock_client, mock_admin_client


















from unittest.mock import MagicMock

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

# unit tests

@pytest.fixture
def mock_astra_db_component():
    component = AstraDBVectorStoreComponent()
    component.token = "mock_token"
    return component

def test_single_database_single_collection(mock_astra_db_component, monkeypatch):
    # Mock the admin client and database client
    mock_admin_client = MagicMock()
    mock_db_client = MagicMock()
    
    # Mock database info
    mock_db = MagicMock()
    mock_db.info.name = "test_db"
    mock_db.info.id = "test_id"
    mock_db.info.region = "test_region"
    
    # Mock list_databases and get_database
    mock_admin_client.list_databases.return_value = [mock_db]
    mock_db_client.list_collection_names.return_value = ["collection1"]
    
    # Patch the DataAPIClient methods
    monkeypatch.setattr(DataAPIClient, "get_admin", lambda self, token: mock_admin_client)
    monkeypatch.setattr(DataAPIClient, "get_database", lambda self, db_id, token: mock_db_client)
    
    codeflash_output = mock_astra_db_component.get_database_list()
    
    expected = {
        "test_db": {
            "api_endpoint": "https://test_id-test_region.apps.astra.datastax.com",
            "collections": 1,
            "records": 0,
        }
    }

def test_multiple_databases_multiple_collections(mock_astra_db_component, monkeypatch):
    # Mock the admin client and database client
    mock_admin_client = MagicMock()
    mock_db_client1 = MagicMock()
    mock_db_client2 = MagicMock()
    
    # Mock database info
    mock_db1 = MagicMock()
    mock_db1.info.name = "test_db1"
    mock_db1.info.id = "test_id1"
    mock_db1.info.region = "test_region1"
    
    mock_db2 = MagicMock()
    mock_db2.info.name = "test_db2"
    mock_db2.info.id = "test_id2"
    mock_db2.info.region = "test_region2"
    
    # Mock list_databases and get_database
    mock_admin_client.list_databases.return_value = [mock_db1, mock_db2]
    mock_db_client1.list_collection_names.return_value = ["collection1", "collection2"]
    mock_db_client2.list_collection_names.return_value = ["collection3"]
    
    # Patch the DataAPIClient methods
    monkeypatch.setattr(DataAPIClient, "get_admin", lambda self, token: mock_admin_client)
    monkeypatch.setattr(DataAPIClient, "get_database", lambda self, db_id, token: mock_db_client1 if db_id == "test_id1" else mock_db_client2)
    
    codeflash_output = mock_astra_db_component.get_database_list()
    
    expected = {
        "test_db1": {
            "api_endpoint": "https://test_id1-test_region1.apps.astra.datastax.com",
            "collections": 2,
            "records": 0,
        },
        "test_db2": {
            "api_endpoint": "https://test_id2-test_region2.apps.astra.datastax.com",
            "collections": 1,
            "records": 0,
        }
    }

def test_no_databases(mock_astra_db_component, monkeypatch):
    # Mock the admin client
    mock_admin_client = MagicMock()
    
    # Mock list_databases
    mock_admin_client.list_databases.return_value = []
    
    # Patch the DataAPIClient methods
    monkeypatch.setattr(DataAPIClient, "get_admin", lambda self, token: mock_admin_client)
    
    codeflash_output = mock_astra_db_component.get_database_list()
    
    expected = {}



def test_special_characters_in_db_info(mock_astra_db_component, monkeypatch):
    # Mock the admin client and database client
    mock_admin_client = MagicMock()
    mock_db_client = MagicMock()
    
    # Mock database info
    mock_db = MagicMock()
    mock_db.info.name = "test_db!@#"
    mock_db.info.id = "test_id!@#"
    mock_db.info.region = "test_region!@#"
    
    # Mock list_databases and get_database
    mock_admin_client.list_databases.return_value = [mock_db]
    mock_db_client.list_collection_names.return_value = ["collection1"]
    
    # Patch the DataAPIClient methods
    monkeypatch.setattr(DataAPIClient, "get_admin", lambda self, token: mock_admin_client)
    monkeypatch.setattr(DataAPIClient, "get_database", lambda self, db_id, token: mock_db_client)
    
    codeflash_output = mock_astra_db_component.get_database_list()
    
    expected = {
        "test_db!@#": {
            "api_endpoint": "https://test_id!@#-test_region!@#.apps.astra.datastax.com",
            "collections": 1,
            "records": 0,
        }
    }
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

:loudspeaker: Feedback on this optimization?

Jan 02 '25 20:01 codeflash-ai[bot]

⚡️ Codeflash found optimizations for this PR

📄 32% (0.32x) speedup for `AstraDBVectorStoreComponent.collection_exists` in `src/backend/base/langflow/components/vectorstores/astradb.py`

⏱️ Runtime : 46.8 milliseconds → 35.4 milliseconds (best of 60 runs)

📝 Explanation and details

Certainly! Here's a more optimized version of the given code.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from unittest.mock import MagicMock, patch

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

# unit tests

@pytest.fixture
def mock_component():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://mock-endpoint"
    component.database_name = "mock_db"
    component.token = "mock_token"
    component.collection_name = "mock_collection"
    component.log = MagicMock()
    return component

def test_collection_exists_valid_collection(mock_component):
    # Mock the DataAPIClient and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["mock_collection", "another_collection"]
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_invalid_collection(mock_component):
    # Mock the DataAPIClient and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["another_collection"]
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_no_api_endpoint(mock_component):
    mock_component.api_endpoint = None
    codeflash_output = mock_component.collection_exists()

def test_collection_exists_no_database_name(mock_component):
    mock_component.database_name = None
    codeflash_output = mock_component.collection_exists()

def test_collection_exists_invalid_token(mock_component):
    mock_component.token = None
    codeflash_output = mock_component.collection_exists()

def test_collection_exists_network_issue(mock_component):
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Network issue")):
        codeflash_output = mock_component.collection_exists()

def test_collection_exists_special_characters_in_name(mock_component):
    mock_component.collection_name = "special_!@#$%^&*()"
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["special_!@#$%^&*()"]
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_long_collection_name(mock_component):
    mock_component.collection_name = "a" * 256
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["a" * 256]
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_empty_collection_name(mock_component):
    mock_component.collection_name = ""
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = [""]
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_large_number_of_collections(mock_component):
    collections = [f"collection_{i}" for i in range(1000)]
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = collections
        mock_get_database.return_value = mock_database

        mock_component.collection_name = "collection_999"
        codeflash_output = mock_component.collection_exists()

def test_collection_exists_exception_logging(mock_component):
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Test exception")):
        codeflash_output = mock_component.collection_exists()
        mock_component.log.assert_called_with("Error getting collection status: Test exception")

def test_collection_exists_case_sensitivity(mock_component):
    mock_component.collection_name = "TestCollection"
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["testcollection"]
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_token_with_special_characters(mock_component):
    mock_component.token = "token_with_special_!@#$%^&*()"
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["mock_collection"]
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_no_collections(mock_component):
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = []
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_identically_named_collections(mock_component):
    mock_component.collection_name = "collection"
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["collection", "Collection"]
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_rate_limiting(mock_component):
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Rate limit exceeded")):
        codeflash_output = mock_component.collection_exists()

def test_collection_exists_schema_changes(mock_component):
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.side_effect = [[], ["mock_collection"]]
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_unexpected_api_response(mock_component):
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = None
        mock_get_database.return_value = mock_database

        codeflash_output = mock_component.collection_exists()

def test_collection_exists_flaky_network(mock_component):
    with patch.object(DataAPIClient, 'get_database', side_effect=[Exception("Network issue"), MagicMock()]):
        codeflash_output = mock_component.collection_exists()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from unittest.mock import MagicMock, patch

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

# unit tests

# Basic Functionality
def test_collection_exists():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = "existing_collection"

    # Mocking the database client and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["existing_collection", "another_collection"]
        mock_get_database.return_value = mock_database
        
        # Assert that collection_exists returns True
        codeflash_output = component.collection_exists()

def test_collection_does_not_exist():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = "non_existing_collection"

    # Mocking the database client and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["existing_collection", "another_collection"]
        mock_get_database.return_value = mock_database
        
        # Assert that collection_exists returns False
        codeflash_output = component.collection_exists()

# API Endpoint Retrieval
def test_api_endpoint_is_set():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"

def test_api_endpoint_not_set_database_name_set():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = "test_db"
    component.get_database_list = MagicMock(return_value={"test_db": "http://fake-endpoint"})

def test_api_endpoint_and_database_name_not_set():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = None

# Token Validity
def test_valid_token():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = "existing_collection"

    # Mocking the database client and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["existing_collection", "another_collection"]
        mock_get_database.return_value = mock_database
        
        # Assert that collection_exists performs as expected
        codeflash_output = component.collection_exists()

def test_invalid_token():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "invalid-token"
    component.collection_name = "existing_collection"

    # Mocking the database client to raise an exception for invalid token
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Authentication failed")):
        # Assert that collection_exists handles authentication failure gracefully
        codeflash_output = component.collection_exists()

# Error Handling
def test_network_error():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = "existing_collection"

    # Mocking the database client to raise a network error
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Network error")):
        # Assert that collection_exists catches the exception and returns False
        codeflash_output = component.collection_exists()

def test_api_endpoint_not_reachable():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://unreachable-endpoint"
    component.token = "valid-token"
    component.collection_name = "existing_collection"

    # Mocking the database client to raise an exception for unreachable endpoint
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Endpoint not reachable")):
        # Assert that collection_exists handles the error and returns False
        codeflash_output = component.collection_exists()

def test_unexpected_exception():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = "existing_collection"

    # Mocking the database client to raise an unexpected exception
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Unexpected error")):
        # Assert that collection_exists catches the exception and returns False
        codeflash_output = component.collection_exists()

# Logging
def test_error_logging():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = "existing_collection"
    component.log = MagicMock()

    # Mocking the database client to raise an exception
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Logging error test")):
        # Call the method
        component.collection_exists()
        
        # Assert that the error message is logged correctly
        component.log.assert_called_with("Error getting collection status: Logging error test")

# Large Scale Test Cases
def test_large_number_of_collections():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = "target_collection"

    # Mocking the database client and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = [f"collection_{i}" for i in range(1000)] + ["target_collection"]
        mock_get_database.return_value = mock_database
        
        # Assert that collection_exists performs efficiently
        codeflash_output = component.collection_exists()

def test_empty_collection_list():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = "non_existing_collection"

    # Mocking the database client and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = []
        mock_get_database.return_value = mock_database
        
        # Assert that collection_exists returns False
        codeflash_output = component.collection_exists()

# Edge Cases
def test_empty_collection_name():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = ""

    # Mocking the database client and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["existing_collection", "another_collection"]
        mock_get_database.return_value = mock_database
        
        # Assert that collection_exists handles it gracefully
        codeflash_output = component.collection_exists()

def test_special_characters_in_collection_name():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = "special!@#collection"

    # Mocking the database client and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["special!@#collection", "another_collection"]
        mock_get_database.return_value = mock_database
        
        # Assert that collection_exists handles it correctly
        codeflash_output = component.collection_exists()

def test_none_collection_name():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = None

    # Mocking the database client and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["existing_collection", "another_collection"]
        mock_get_database.return_value = mock_database
        
        # Assert that collection_exists handles it gracefully
        codeflash_output = component.collection_exists()

# Side Effects
def test_state_mutation():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "http://fake-endpoint"
    component.token = "valid-token"
    component.collection_name = "existing_collection"
    initial_state = component.__dict__.copy()

    # Mocking the database client and its methods
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["existing_collection", "another_collection"]
        mock_get_database.return_value = mock_database
        
        # Call the method
        component.collection_exists()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

:loudspeaker: Feedback on this optimization?

Jan 02 '25 20:01 codeflash-ai[bot]

⚡️ Codeflash found optimizations for this PR

📄 11% (0.11x) speedup for `AstraDBVectorStoreComponent.collection_exists` in `src/backend/base/langflow/components/vectorstores/astradb.py`

⏱️ Runtime : 2.48 milliseconds → 2.23 milliseconds (best of 81 runs)

📝 Explanation and details

Certainly! Here is an optimized version of the program.

Key improvements.

Added an early exit condition in collection_exists to return False if required attributes (api_endpoint, collection_name) are missing to avoid unnecessary API calls.
Removed unnecessary list() conversion as in operator works directly on the collections returned by list_collections(), assuming it returns an iterable or set-like object.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from unittest.mock import MagicMock, patch

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

# unit tests

# Helper function to create a mock component
def create_mock_component(api_endpoint=None, database_name=None, token=None, collection_name=None):
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = api_endpoint
    component.database_name = database_name
    component.token = token
    component.collection_name = collection_name
    component.log = MagicMock()
    component.get_database_list = MagicMock(return_value={database_name: api_endpoint})
    return component

# Basic Functionality
def test_collection_exists():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="valid_token", collection_name="test_collection")
    
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collections=MagicMock(return_value=["test_collection"]))):
        codeflash_output = component.collection_exists()

def test_collection_does_not_exist():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="valid_token", collection_name="non_existent_collection")
    
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collections=MagicMock(return_value=["test_collection"]))):
        codeflash_output = component.collection_exists()

# Uninitialized Attributes
def test_uninitialized_api_endpoint():
    component = create_mock_component(api_endpoint=None, database_name="test_db", token="valid_token", collection_name="test_collection")
    codeflash_output = component.collection_exists()

def test_uninitialized_database_name():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name=None, token="valid_token", collection_name="test_collection")
    codeflash_output = component.collection_exists()

def test_uninitialized_token():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token=None, collection_name="test_collection")
    codeflash_output = component.collection_exists()

def test_uninitialized_collection_name():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="valid_token", collection_name=None)
    codeflash_output = component.collection_exists()

# Error Handling
def test_network_error():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="valid_token", collection_name="test_collection")
    
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Network error")):
        codeflash_output = component.collection_exists()
        component.log.assert_called_with("Error getting collection status: Network error")

def test_authentication_error():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="invalid_token", collection_name="test_collection")
    
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Authentication error")):
        codeflash_output = component.collection_exists()
        component.log.assert_called_with("Error getting collection status: Authentication error")

def test_api_endpoint_not_reachable():
    component = create_mock_component(api_endpoint="http://invalid_endpoint", database_name="test_db", token="valid_token", collection_name="test_collection")
    
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Endpoint not reachable")):
        codeflash_output = component.collection_exists()
        component.log.assert_called_with("Error getting collection status: Endpoint not reachable")

def test_database_not_found():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="non_existent_db", token="valid_token", collection_name="test_collection")
    
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Database not found")):
        codeflash_output = component.collection_exists()
        component.log.assert_called_with("Error getting collection status: Database not found")

# Edge Cases
def test_empty_collection_list():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="valid_token", collection_name="test_collection")
    
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collections=MagicMock(return_value=[]))):
        codeflash_output = component.collection_exists()

def test_large_number_of_collections():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="valid_token", collection_name="test_collection")
    large_collection_list = [f"collection_{i}" for i in range(1000)] + ["test_collection"]
    
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collections=MagicMock(return_value=large_collection_list))):
        codeflash_output = component.collection_exists()

def test_special_characters_in_collection_name():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="valid_token", collection_name="special!@#collection")
    
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collections=MagicMock(return_value=["special!@#collection"]))):
        codeflash_output = component.collection_exists()

def test_case_sensitivity():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="valid_token", collection_name="Collection_Name")
    
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collections=MagicMock(return_value=["collection_name", "Collection_Name"]))):
        codeflash_output = component.collection_exists()

# Performance and Scalability
def test_performance_large_scale():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="valid_token", collection_name="test_collection")
    large_collection_list = [f"collection_{i}" for i in range(1000)] + ["test_collection"]
    
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collections=MagicMock(return_value=large_collection_list))):
        codeflash_output = component.collection_exists()

# Side Effects
def test_logging_on_error():
    component = create_mock_component(api_endpoint="http://valid_endpoint", database_name="test_db", token="valid_token", collection_name="test_collection")
    
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Test error")):
        codeflash_output = component.collection_exists()
        component.log.assert_called_with("Error getting collection status: Test error")
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from unittest.mock import MagicMock, patch

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

# unit tests

@pytest.fixture
def mock_component():
    component = AstraDBVectorStoreComponent()
    component.token = "valid_token"
    component.api_endpoint = "http://valid.endpoint"
    component.database_name = "test_database"
    component.collection_name = "test_collection"
    component.get_database_list = MagicMock(return_value={"test_database": "http://valid.endpoint"})
    component.log = MagicMock()
    return component

def test_collection_exists(mock_component):
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["test_collection", "another_collection"]
        mock_get_database.return_value = mock_database
        
        codeflash_output = mock_component.collection_exists()

def test_collection_does_not_exist(mock_component):
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["another_collection"]
        mock_get_database.return_value = mock_database
        
        codeflash_output = mock_component.collection_exists()

def test_api_endpoint_set(mock_component):
    mock_component.api_endpoint = "http://explicit.endpoint"

def test_api_endpoint_from_database_name(mock_component):
    mock_component.api_endpoint = None

def test_no_api_endpoint_or_database_name(mock_component):
    mock_component.api_endpoint = None
    mock_component.database_name = None

def test_invalid_token(mock_component):
    mock_component.token = "invalid_token"
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Invalid token")):
        codeflash_output = mock_component.collection_exists()
        mock_component.log.assert_called_with("Error getting collection status: Invalid token")

def test_network_issues(mock_component):
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Network error")):
        codeflash_output = mock_component.collection_exists()
        mock_component.log.assert_called_with("Error getting collection status: Network error")

def test_invalid_api_endpoint(mock_component):
    mock_component.api_endpoint = "http://invalid.endpoint"
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Invalid endpoint")):
        codeflash_output = mock_component.collection_exists()
        mock_component.log.assert_called_with("Error getting collection status: Invalid endpoint")

def test_database_not_found(mock_component):
    mock_component.database_name = "non_existent_database"
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Database not found")):
        codeflash_output = mock_component.collection_exists()
        mock_component.log.assert_called_with("Error getting collection status: Database not found")

def test_empty_collection_name(mock_component):
    mock_component.collection_name = ""
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["test_collection"]
        mock_get_database.return_value = mock_database
        
        codeflash_output = mock_component.collection_exists()

def test_special_characters_in_collection_name(mock_component):
    mock_component.collection_name = "special@collection!"
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["special@collection!", "another_collection"]
        mock_get_database.return_value = mock_database
        
        codeflash_output = mock_component.collection_exists()

def test_large_number_of_collections(mock_component):
    collections = [f"collection_{i}" for i in range(1000)]
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = collections
        mock_get_database.return_value = mock_database
        
        mock_component.collection_name = "collection_999"
        codeflash_output = mock_component.collection_exists()

def test_case_sensitivity(mock_component):
    mock_component.collection_name = "TestCollection"
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = ["testcollection"]
        mock_get_database.return_value = mock_database
        
        codeflash_output = mock_component.collection_exists()

def test_logging_on_error(mock_component):
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Some error")):
        codeflash_output = mock_component.collection_exists()
        mock_component.log.assert_called_with("Error getting collection status: Some error")

def test_performance_with_large_data(mock_component):
    collections = [f"collection_{i}" for i in range(1000)]
    with patch.object(DataAPIClient, 'get_database') as mock_get_database:
        mock_database = MagicMock()
        mock_database.list_collections.return_value = collections
        mock_get_database.return_value = mock_database
        
        mock_component.collection_name = "collection_500"
        codeflash_output = mock_component.collection_exists()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

:loudspeaker: Feedback on this optimization?

Jan 03 '25 16:01 codeflash-ai[bot]

⚡️ Codeflash found optimizations for this PR

📄 17% (0.17x) speedup for `AstraDBVectorStoreComponent.get_api_endpoint` in `src/backend/base/langflow/components/vectorstores/astradb.py`

⏱️ Runtime : 27.6 milliseconds → 23.6 milliseconds (best of 18 runs)

📝 Explanation and details

Changes made to optimize performance.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from unittest.mock import MagicMock, patch

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

# unit tests

# Basic Functionality
def test_api_endpoint_already_set():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "https://valid.url"
    codeflash_output = component.get_api_endpoint()

def test_database_name_set_existing():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = "existing_db"
    component.token = "valid_token"

    mock_db = MagicMock()
    mock_db.info.name = "existing_db"
    mock_db.info.id = "dbid"
    mock_db.info.region = "region"

    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db]))):
        with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["col1", "col2"]))):
            codeflash_output = component.get_api_endpoint()


def test_token_validity():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = "existing_db"
    component.token = "valid_token"

    mock_db = MagicMock()
    mock_db.info.name = "existing_db"
    mock_db.info.id = "dbid"
    mock_db.info.region = "region"

    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db]))):
        with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["col1", "col2"]))):
            codeflash_output = component.get_api_endpoint()

def test_empty_database_list():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = "existing_db"
    component.token = "valid_token"

    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[]))):
        codeflash_output = component.get_api_endpoint()

def test_database_name_not_set():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = None
    codeflash_output = component.get_api_endpoint()

def test_database_name_empty_string():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = ""
    codeflash_output = component.get_api_endpoint()

# Error Handling

def test_admin_client_initialization_error():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = "existing_db"
    component.token = "valid_token"

    with patch.object(DataAPIClient, 'get_admin', side_effect=Exception("Admin client error")):
        with pytest.raises(Exception, match="Admin client error"):
            component.get_api_endpoint()

def test_database_list_retrieval_error():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = "existing_db"
    component.token = "valid_token"

    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(side_effect=Exception("Database list error")))):
        with pytest.raises(Exception, match="Database list error"):
            component.get_api_endpoint()

def test_collection_names_retrieval_error():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = "existing_db"
    component.token = "valid_token"

    mock_db = MagicMock()
    mock_db.info.name = "existing_db"
    mock_db.info.id = "dbid"
    mock_db.info.region = "region"

    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db]))):
        with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(side_effect=Exception("Collection names error")))):
            with pytest.raises(Exception, match="Collection names error"):
                component.get_api_endpoint()

# Performance and Scalability
def test_large_number_of_databases():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = "db500"
    component.token = "valid_token"

    mock_dbs = []
    for i in range(1000):
        mock_db = MagicMock()
        mock_db.info.name = f"db{i}"
        mock_db.info.id = f"id{i}"
        mock_db.info.region = "region"
        mock_dbs.append(mock_db)

    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=mock_dbs))):
        with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["col1", "col2"]))):
            codeflash_output = component.get_api_endpoint()

def test_large_number_of_collections():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = None
    component.database_name = "existing_db"
    component.token = "valid_token"

    mock_db = MagicMock()
    mock_db.info.name = "existing_db"
    mock_db.info.id = "dbid"
    mock_db.info.region = "region"

    collections = [f"col{i}" for i in range(1000)]

    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db]))):
        with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=collections))):
            codeflash_output = component.get_api_endpoint()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from unittest.mock import MagicMock, patch

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

# unit tests

# Basic Functionality

def test_api_endpoint_already_set():
    component = AstraDBVectorStoreComponent()
    component.api_endpoint = "https://example.com"
    codeflash_output = component.get_api_endpoint()

def test_database_name_not_set():
    component = AstraDBVectorStoreComponent()
    component.database_name = None
    codeflash_output = component.get_api_endpoint()

def test_valid_database_name():
    component = AstraDBVectorStoreComponent()
    component.database_name = "valid_db_name"
    mock_db_list = {
        "valid_db_name": {
            "api_endpoint": "https://valid-db-id-region.apps.astra.datastax.com",
            "collections": 5,
            "records": 0,
        }
    }
    component.get_database_list = MagicMock(return_value=mock_db_list)
    codeflash_output = component.get_api_endpoint()

def test_invalid_database_name():
    component = AstraDBVectorStoreComponent()
    component.database_name = "invalid_db_name"
    mock_db_list = {
        "valid_db_name": {
            "api_endpoint": "https://valid-db-id-region.apps.astra.datastax.com",
            "collections": 5,
            "records": 0,
        }
    }
    component.get_database_list = MagicMock(return_value=mock_db_list)
    codeflash_output = component.get_api_endpoint()

# Edge Cases

def test_token_is_none():
    component = AstraDBVectorStoreComponent()
    component.token = None
    with pytest.raises(Exception):
        component.get_database_list()

def test_empty_database_list():
    component = AstraDBVectorStoreComponent()
    component.get_database_list = MagicMock(return_value={})
    component.database_name = "any_db_name"
    codeflash_output = component.get_api_endpoint()

def test_database_with_no_collections():
    component = AstraDBVectorStoreComponent()
    component.database_name = "db_with_no_collections"
    mock_db_list = {
        "db_with_no_collections": {
            "api_endpoint": "https://db-id-region.apps.astra.datastax.com",
            "collections": 0,
            "records": 0,
        }
    }
    component.get_database_list = MagicMock(return_value=mock_db_list)
    codeflash_output = component.get_api_endpoint()

# Error Handling

def test_data_api_client_initialization_failure():
    with patch('astrapy.DataAPIClient') as MockDataAPIClient:
        MockDataAPIClient.side_effect = Exception("Initialization failed")
        component = AstraDBVectorStoreComponent()
        component.token = "valid_token"
        with pytest.raises(Exception):
            component.get_database_list()

def test_admin_client_get_admin_failure():
    with patch('astrapy.DataAPIClient.get_admin') as MockGetAdmin:
        MockGetAdmin.side_effect = Exception("Get admin failed")
        component = AstraDBVectorStoreComponent()
        component.token = "valid_token"
        with pytest.raises(Exception):
            component.get_database_list()

def test_list_databases_failure():
    with patch('astrapy.DataAPIClient.get_admin') as MockGetAdmin:
        mock_admin = MockGetAdmin.return_value
        mock_admin.list_databases.side_effect = Exception("List databases failed")
        component = AstraDBVectorStoreComponent()
        component.token = "valid_token"
        with pytest.raises(Exception):
            component.get_database_list()

def test_list_collection_names_failure():
    with patch('astrapy.DataAPIClient') as MockDataAPIClient:
        mock_client = MockDataAPIClient.return_value
        mock_admin = mock_client.get_admin.return_value
        mock_admin.list_databases.return_value = [MagicMock(info=MagicMock(name="db_name", id="db_id", region="region"))]
        mock_client.get_database.return_value.list_collection_names.side_effect = Exception("List collections failed")
        component = AstraDBVectorStoreComponent()
        component.token = "valid_token"
        with pytest.raises(Exception):
            component.get_database_list()

# Performance and Scalability

def test_large_number_of_databases():
    component = AstraDBVectorStoreComponent()
    component.database_name = "db_500"
    mock_db_list = {
        f"db_{i}": {
            "api_endpoint": f"https://db-{i}-region.apps.astra.datastax.com",
            "collections": 5,
            "records": 0,
        }
        for i in range(1000)
    }
    component.get_database_list = MagicMock(return_value=mock_db_list)
    codeflash_output = component.get_api_endpoint()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

:loudspeaker: Feedback on this optimization?

Jan 06 '25 21:01 codeflash-ai[bot]

⚡️ Codeflash found optimizations for this PR

📄 11% (0.11x) speedup for `AstraDBVectorStoreComponent.get_database_list` in `src/backend/base/langflow/components/vectorstores/astradb.py`

⏱️ Runtime : 32.0 milliseconds → 28.9 milliseconds (best of 5 runs)

📝 Explanation and details

Certainly! Here is an optimized version of the code with the aim to minimize repeated lookups and function calls.

The optimized code ensures that we avoid unnecessary list conversions and redundant API calls by storing the required values in variables and reusing them where necessary.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from unittest.mock import MagicMock, patch

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

# unit tests

# Helper function to create mock database objects
def create_mock_db(name, id, region, collections):
    db = MagicMock()
    db.info.name = name
    db.info.id = id
    db.info.region = region
    db.list_collection_names = MagicMock(return_value=collections)
    return db

# Basic Functionality Tests

def test_multiple_databases_with_collections():
    component = AstraDBVectorStoreComponent()
    component.token = "valid_token"
    
    mock_db1 = create_mock_db("db1", "db1-id", "us-west", ["col1", "col2", "col3"])
    mock_db2 = create_mock_db("db2", "db2-id", "us-east", ["col1", "col2", "col3", "col4", "col5"])
    
    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db1, mock_db2]))):
        with patch.object(DataAPIClient, 'get_database', side_effect=[mock_db1, mock_db2]):
            codeflash_output = component.get_database_list()
            expected = {
                "db1": {
                    "api_endpoint": "https://db1-id-us-west.apps.astra.datastax.com",
                    "collections": 3,
                    "records": 0
                },
                "db2": {
                    "api_endpoint": "https://db2-id-us-east.apps.astra.datastax.com",
                    "collections": 5,
                    "records": 0
                }
            }

# Edge Case Tests


def test_database_with_long_name():
    component = AstraDBVectorStoreComponent()
    component.token = "valid_token"
    
    long_name = "a" * 256
    mock_db = create_mock_db(long_name, "long-name-id", "us-west", [])
    
    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db]))):
        with patch.object(DataAPIClient, 'get_database', return_value=mock_db):
            codeflash_output = component.get_database_list()
            expected = {
                long_name: {
                    "api_endpoint": "https://long-name-id-us-west.apps.astra.datastax.com",
                    "collections": 0,
                    "records": 0
                }
            }

# Error Handling Tests
def test_invalid_token():
    component = AstraDBVectorStoreComponent()
    component.token = "invalid_token"
    
    with patch.object(DataAPIClient, 'get_admin', side_effect=Exception("Invalid token")):
        with pytest.raises(Exception, match="Invalid token"):
            component.get_database_list()

def test_api_failure():
    component = AstraDBVectorStoreComponent()
    component.token = "valid_token"
    
    with patch.object(DataAPIClient, 'get_admin', side_effect=Exception("API failure")):
        with pytest.raises(Exception, match="API failure"):
            component.get_database_list()

def test_partial_data_retrieval():
    component = AstraDBVectorStoreComponent()
    component.token = "valid_token"
    
    mock_db1 = create_mock_db("db1", "db1-id", "us-west", ["col1", "col2"])
    
    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db1]))):
        with patch.object(DataAPIClient, 'get_database', side_effect=[mock_db1, Exception("API failure")]):
            codeflash_output = component.get_database_list()
            expected = {
                "db1": {
                    "api_endpoint": "https://db1-id-us-west.apps.astra.datastax.com",
                    "collections": 2,
                    "records": 0
                }
            }

# Performance and Scalability Tests


def test_verify_api_calls():
    component = AstraDBVectorStoreComponent()
    component.token = "valid_token"
    
    mock_db = create_mock_db("db1", "db1-id", "us-west", [])
    
    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db]))) as mock_get_admin:
        with patch.object(DataAPIClient, 'get_database', return_value=mock_db) as mock_get_database:
            component.get_database_list()
            mock_get_admin.assert_called_once_with(token="valid_token")
            mock_get_database.assert_called_once_with(api_endpoint="https://db1-id-us-west.apps.astra.datastax.com", token="valid_token")

def test_token_usage():
    component = AstraDBVectorStoreComponent()
    component.token = "valid_token"
    
    mock_db = create_mock_db("db1", "db1-id", "us-west", [])
    
    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db]))) as mock_get_admin:
        with patch.object(DataAPIClient, 'get_database', return_value=mock_db) as mock_get_database:
            component.get_database_list()
            mock_get_admin.assert_called_once_with(token="valid_token")
            mock_get_database.assert_called_once_with(api_endpoint="https://db1-id-us-west.apps.astra.datastax.com", token="valid_token")

# Special Scenarios Tests
def test_database_in_multiple_regions():
    component = AstraDBVectorStoreComponent()
    component.token = "valid_token"
    
    mock_db1 = create_mock_db("db1", "db1-id", "us-west", [])
    mock_db2 = create_mock_db("db1", "db1-id", "us-east", [])
    
    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db1, mock_db2]))):
        with patch.object(DataAPIClient, 'get_database', side_effect=[mock_db1, mock_db2]):
            codeflash_output = component.get_database_list()
            expected = {
                "db1": {
                    "api_endpoint": "https://db1-id-us-west.apps.astra.datastax.com",
                    "collections": 0,
                    "records": 0
                },
                "db1": {
                    "api_endpoint": "https://db1-id-us-east.apps.astra.datastax.com",
                    "collections": 0,
                    "records": 0
                }
            }

def test_concurrent_execution():
    component = AstraDBVectorStoreComponent()
    component.token = "valid_token"
    
    mock_db = create_mock_db("db1", "db1-id", "us-west", [])
    
    with patch.object(DataAPIClient, 'get_admin', return_value=MagicMock(list_databases=MagicMock(return_value=[mock_db]))):
        with patch.object(DataAPIClient, 'get_database', return_value=mock_db):
            import concurrent.futures
            with concurrent.futures.ThreadPoolExecutor() as executor:
                futures = [executor.submit(component.get_database_list) for _ in range(10)]
                results = [future.result() for future in futures]
                expected = {
                    "db1": {
                        "api_endpoint": "https://db1-id-us-west.apps.astra.datastax.com",
                        "collections": 0,
                        "records": 0
                    }
                }
                for result in results:
                    pass
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from unittest.mock import MagicMock, patch

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

# unit tests

@pytest.fixture
def mock_data_api_client():
    with patch('astrapy.DataAPIClient') as MockClient:
        yield MockClient

@pytest.fixture
def mock_astra_db_vector_store_component(mock_data_api_client):
    component = AstraDBVectorStoreComponent()
    component.token = 'fake-token'
    return component









def test_api_endpoint_construction_failure(mock_astra_db_vector_store_component, mock_data_api_client):
    # Mocking the responses
    mock_admin_client = MagicMock()
    mock_admin_client.list_databases.return_value = [{'info': {'name': 'db1', 'id': '', 'region': 'us-west'}}]
    mock_data_api_client.return_value.get_admin.return_value = mock_admin_client

    with pytest.raises(ValueError):
        mock_astra_db_vector_store_component.get_database_list()

:loudspeaker: Feedback on this optimization?

Jan 06 '25 22:01 codeflash-ai[bot]

⚡️ Codeflash found optimizations for this PR

📄 17% (0.17x) speedup for `AstraDBVectorStoreComponent.collection_exists` in `src/backend/base/langflow/components/vectorstores/astradb.py`

⏱️ Runtime : 6.54 milliseconds → 5.61 milliseconds (best of 55 runs)

📝 Explanation and details

Sure, here is a faster version of your code.

Summary of Improvements.

Condition Optimization in get_keyspace.
- Combined the condition checking and stripping into one line for efficiency.
Early Return in get_api_endpoint.
- Simplified the method to reduce unnecessary checks and nesting.
- Directly chained .get methods with default values to prevent key errors.
Early Exit in collection_exists.
- Added checks for all necessary instance variables at the start to prevent unnecessary API calls.
- Cached the results of get_api_endpoint and get_keyspace function calls for reuse.
- Reduced the number of times get_keyspace is called by saving its result.

These changes should help the code run faster by reducing redundant checks and minimizing the number of method calls.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from unittest.mock import MagicMock, patch

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


# unit tests
class TestAstraDBVectorStoreComponent:
    @pytest.fixture
    def setup_component(self):
        component = AstraDBVectorStoreComponent()
        component.token = "valid_token"
        component.keyspace = "valid_keyspace"
        component.api_endpoint = "http://valid.endpoint"
        component.database_name = "valid_database"
        component.collection_name = "valid_collection"
        component.get_database_list = MagicMock(return_value={"valid_database": {"api_endpoint": "http://valid.endpoint"}})
        component.log = MagicMock()
        return component

    def test_collection_exists(self, setup_component):
        component = setup_component
        mock_database = MagicMock()
        mock_database.list_collection_names.return_value = ["valid_collection", "another_collection"]
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            codeflash_output = component.collection_exists()

    def test_collection_does_not_exist(self, setup_component):
        component = setup_component
        mock_database = MagicMock()
        mock_database.list_collection_names.return_value = ["another_collection"]
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            codeflash_output = component.collection_exists()

    def test_empty_keyspace(self, setup_component):
        component = setup_component
        component.keyspace = " "
        
        codeflash_output = component.collection_exists()

    def test_empty_api_endpoint(self, setup_component):
        component = setup_component
        component.api_endpoint = ""
        
        codeflash_output = component.collection_exists()

    def test_empty_database_name(self, setup_component):
        component = setup_component
        component.database_name = ""
        
        codeflash_output = component.collection_exists()

    def test_empty_collection_name(self, setup_component):
        component = setup_component
        component.collection_name = ""
        
        codeflash_output = component.collection_exists()

    def test_invalid_keyspace(self, setup_component):
        component = setup_component
        component.keyspace = "invalid keyspace"
        mock_database = MagicMock()
        mock_database.list_collection_names.side_effect = Exception("Invalid keyspace")
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            codeflash_output = component.collection_exists()

    def test_invalid_api_endpoint(self, setup_component):
        component = setup_component
        component.api_endpoint = "http://invalid.endpoint"
        mock_database = MagicMock()
        mock_database.list_collection_names.side_effect = Exception("Invalid API endpoint")
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            codeflash_output = component.collection_exists()

    def test_invalid_token(self, setup_component):
        component = setup_component
        component.token = "invalid_token"
        mock_database = MagicMock()
        mock_database.list_collection_names.side_effect = Exception("Invalid token")
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            codeflash_output = component.collection_exists()

    def test_api_client_initialization_failure(self, setup_component):
        component = setup_component
        
        with patch('astrapy.DataAPIClient', side_effect=Exception("Initialization failure")):
            codeflash_output = component.collection_exists()

    def test_database_retrieval_failure(self, setup_component):
        component = setup_component
        
        with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Database retrieval failure")):
            codeflash_output = component.collection_exists()

    def test_collection_list_retrieval_failure(self, setup_component):
        component = setup_component
        mock_database = MagicMock()
        mock_database.list_collection_names.side_effect = Exception("Collection list retrieval failure")
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            codeflash_output = component.collection_exists()

    def test_logging_side_effect(self, setup_component):
        component = setup_component
        mock_database = MagicMock()
        mock_database.list_collection_names.side_effect = Exception("Logging test")
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            component.collection_exists()
            component.log.assert_called_with("Error getting collection status: Logging test")

    def test_large_number_of_collections(self, setup_component):
        component = setup_component
        collections = [f"collection_{i}" for i in range(1000)]
        mock_database = MagicMock()
        mock_database.list_collection_names.return_value = collections
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            codeflash_output = component.collection_exists()

    def test_large_collection_names(self, setup_component):
        component = setup_component
        large_name = "a" * 1000
        component.collection_name = large_name
        mock_database = MagicMock()
        mock_database.list_collection_names.return_value = [large_name]
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            codeflash_output = component.collection_exists()

    def test_special_characters_in_collection_names(self, setup_component):
        component = setup_component
        special_name = "colléction_ñame"
        component.collection_name = special_name
        mock_database = MagicMock()
        mock_database.list_collection_names.return_value = [special_name]
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            codeflash_output = component.collection_exists()

    def test_performance_with_large_data_samples(self, setup_component):
        component = setup_component
        collections = [f"collection_{i}" for i in range(1000)]
        mock_database = MagicMock()
        mock_database.list_collection_names.return_value = collections
        
        with patch.object(DataAPIClient, 'get_database', return_value=mock_database):
            codeflash_output = component.collection_exists()

    

from unittest.mock import MagicMock, patch

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

# unit tests

# Helper function to create a mock AstraDBVectorStoreComponent
def create_mock_component(keyspace, collection_name, api_endpoint=None, database_name=None, token=None):
    component = AstraDBVectorStoreComponent()
    component.keyspace = keyspace
    component.collection_name = collection_name
    component.api_endpoint = api_endpoint
    component.database_name = database_name
    component.token = token
    component.get_database_list = MagicMock(return_value={database_name: {"api_endpoint": api_endpoint}})
    component.log = MagicMock()
    return component

# Normal Operation
def test_collection_exists_normal():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["existing_collection"]))):
        codeflash_output = component.collection_exists()

def test_collection_not_exists_normal():
    component = create_mock_component("valid_keyspace", "non_existing_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["existing_collection"]))):
        codeflash_output = component.collection_exists()

# Edge Cases
def test_empty_keyspace():
    component = create_mock_component("", "any_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    codeflash_output = component.collection_exists()

def test_whitespace_keyspace():
    component = create_mock_component("    ", "any_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    codeflash_output = component.collection_exists()

def test_empty_collection_name():
    component = create_mock_component("valid_keyspace", "", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["existing_collection"]))):
        codeflash_output = component.collection_exists()

def test_whitespace_collection_name():
    component = create_mock_component("valid_keyspace", "    ", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["existing_collection"]))):
        codeflash_output = component.collection_exists()

# Invalid Inputs
def test_invalid_keyspace_format():
    component = create_mock_component("invalid!keyspace", "any_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["existing_collection"]))):
        codeflash_output = component.collection_exists()

def test_invalid_collection_name_format():
    component = create_mock_component("valid_keyspace", "invalid!collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["existing_collection"]))):
        codeflash_output = component.collection_exists()

# API Endpoint Scenarios
def test_valid_api_endpoint():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["existing_collection"]))):
        codeflash_output = component.collection_exists()

def test_missing_api_endpoint():
    component = create_mock_component("valid_keyspace", "existing_collection", None, "valid_database", "valid_token")
    codeflash_output = component.collection_exists()

def test_invalid_api_endpoint():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://invalid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Invalid API endpoint")):
        codeflash_output = component.collection_exists()

# Authentication Scenarios
def test_valid_token():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["existing_collection"]))):
        codeflash_output = component.collection_exists()

def test_invalid_token():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", "invalid_token")
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Invalid token")):
        codeflash_output = component.collection_exists()

def test_missing_token():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", None)
    codeflash_output = component.collection_exists()

# Database Connection Scenarios
def test_database_connection_successful():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["existing_collection"]))):
        codeflash_output = component.collection_exists()

def test_database_connection_fails():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "invalid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Database connection failed")):
        codeflash_output = component.collection_exists()

# Exception Handling
def test_api_client_throws_exception():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, '__init__', side_effect=Exception("API client exception")):
        codeflash_output = component.collection_exists()

def test_database_retrieval_throws_exception():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Database retrieval exception")):
        codeflash_output = component.collection_exists()

def test_collection_listing_throws_exception():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(side_effect=Exception("Collection listing exception")))):
        codeflash_output = component.collection_exists()

# Performance and Scalability
def test_large_number_of_collections():
    collections = [f"collection_{i}" for i in range(1000)]
    component = create_mock_component("valid_keyspace", "collection_999", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=collections))):
        codeflash_output = component.collection_exists()

def test_high_frequency_of_calls():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["existing_collection"]))):
        for _ in range(100):
            codeflash_output = component.collection_exists()

# Boundary Conditions
def test_minimum_length_keyspace_and_collection_names():
    component = create_mock_component("k", "c", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=["c"]))):
        codeflash_output = component.collection_exists()

def test_maximum_length_keyspace_and_collection_names():
    keyspace = "k" * 255
    collection = "c" * 255
    component = create_mock_component(keyspace, collection, "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', return_value=MagicMock(list_collection_names=MagicMock(return_value=[collection]))):
        codeflash_output = component.collection_exists()

# Side Effects Verification
def test_logging_on_exception():
    component = create_mock_component("valid_keyspace", "existing_collection", "https://valid.api.endpoint", "valid_database", "valid_token")
    with patch.object(DataAPIClient, 'get_database', side_effect=Exception("Test exception")):
        codeflash_output = component.collection_exists()
        component.log.assert_called_with("Error getting collection status: Test exception")
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

:loudspeaker: Feedback on this optimization?

Jan 06 '25 22:01 codeflash-ai[bot]

@ogabrielluiz the only reason i'm thinking the astra DB changes should be combined with this PR is that it brings into sync the "secret sauce" (for lack of a better term) functionality that we had in DS LF... in other words, if they are separated, i dont think the functionality will work at all.... if we merge in the UI side dialog inputs, the creation logic for a new DB or collection in the OSS version of the component would be redundant, and if we merged in the Astra DB side changes, then there'd be no way to create a collection.

That's my brief argument for making keeping them together but let me know if you want to chat about it :)

FYI this is still a WIP but we're getting very close and would love a second look from you when you can.

Jan 07 '25 20:01 erichare

@ogabrielluiz i think the PR title has been updated, do things look good here now to you? Was wondering if we could get an approval to test the full CI (feel free to mark as Do Not Merge if you want to wait on the actual merge, since obviously there is a pretty core new addition here).

Jan 09 '25 19:01 erichare

feat: add dialog support to dropdown components and enhance AstraDB i…

CodSpeed Performance Report

Merging #5349 will degrade performances by 91.11%

Summary

Benchmarks breakdown

⚡️ Codeflash found optimizations for this PR

📄 27% (0.27x) speedup for AstraDBVectorStoreComponent.get_database_list in src/backend/base/langflow/components/vectorstores/astradb.py

⚡️ Codeflash found optimizations for this PR

📄 32% (0.32x) speedup for AstraDBVectorStoreComponent.collection_exists in src/backend/base/langflow/components/vectorstores/astradb.py

⚡️ Codeflash found optimizations for this PR

📄 11% (0.11x) speedup for AstraDBVectorStoreComponent.collection_exists in src/backend/base/langflow/components/vectorstores/astradb.py

⚡️ Codeflash found optimizations for this PR

📄 17% (0.17x) speedup for AstraDBVectorStoreComponent.get_api_endpoint in src/backend/base/langflow/components/vectorstores/astradb.py

⚡️ Codeflash found optimizations for this PR

📄 11% (0.11x) speedup for AstraDBVectorStoreComponent.get_database_list in src/backend/base/langflow/components/vectorstores/astradb.py

⚡️ Codeflash found optimizations for this PR

📄 17% (0.17x) speedup for AstraDBVectorStoreComponent.collection_exists in src/backend/base/langflow/components/vectorstores/astradb.py

Summary of Improvements.

📄 27% (0.27x) speedup for `AstraDBVectorStoreComponent.get_database_list` in `src/backend/base/langflow/components/vectorstores/astradb.py`

📄 32% (0.32x) speedup for `AstraDBVectorStoreComponent.collection_exists` in `src/backend/base/langflow/components/vectorstores/astradb.py`

📄 11% (0.11x) speedup for `AstraDBVectorStoreComponent.collection_exists` in `src/backend/base/langflow/components/vectorstores/astradb.py`

📄 17% (0.17x) speedup for `AstraDBVectorStoreComponent.get_api_endpoint` in `src/backend/base/langflow/components/vectorstores/astradb.py`

📄 11% (0.11x) speedup for `AstraDBVectorStoreComponent.get_database_list` in `src/backend/base/langflow/components/vectorstores/astradb.py`

📄 17% (0.17x) speedup for `AstraDBVectorStoreComponent.collection_exists` in `src/backend/base/langflow/components/vectorstores/astradb.py`