mem0ai--mem0/tests/vector_stores/test_azure_mysql.py

import json
from unittest.mock import Mock, patch

import pytest
from pydantic import ValidationError

from mem0.configs.vector_stores.azure_mysql import AzureMySQLConfig
from mem0.vector_stores.azure_mysql import AzureMySQL, OutputData


@pytest.fixture
def mock_connection_pool():
    """Create a mock connection pool."""
    pool = Mock()
    conn = Mock()
    cursor = Mock()

    # Setup cursor mock
    cursor.fetchall = Mock(return_value=[])
    cursor.fetchone = Mock(return_value=None)
    cursor.execute = Mock()
    cursor.executemany = Mock()
    cursor.close = Mock()

    # Setup connection mock
    conn.cursor = Mock(return_value=cursor)
    conn.commit = Mock()
    conn.rollback = Mock()
    conn.close = Mock()

    # Setup pool mock
    pool.connection = Mock(return_value=conn)
    pool.close = Mock()

    return pool


@pytest.fixture
def azure_mysql_instance(mock_connection_pool):
    """Create an AzureMySQL instance with mocked connection pool."""
    with patch('mem0.vector_stores.azure_mysql.PooledDB') as mock_pooled_db:
        mock_pooled_db.return_value = mock_connection_pool

        instance = AzureMySQL(
            host="test-server.mysql.database.azure.com",
            port=3306,
            user="testuser",
            password="testpass",
            database="testdb",
            collection_name="test_collection",
            embedding_model_dims=128,
            use_azure_credential=False,
            ssl_disabled=True,
        )
        instance.connection_pool = mock_connection_pool
        return instance


def test_azure_mysql_init(mock_connection_pool):
    """Test AzureMySQL initialization."""
    with patch('mem0.vector_stores.azure_mysql.PooledDB') as mock_pooled_db:
        mock_pooled_db.return_value = mock_connection_pool

        instance = AzureMySQL(
            host="test-server.mysql.database.azure.com",
            port=3306,
            user="testuser",
            password="testpass",
            database="testdb",
            collection_name="test_collection",
            embedding_model_dims=128,
        )

        assert instance.host == "test-server.mysql.database.azure.com"
        assert instance.port == 3306
        assert instance.user == "testuser"
        assert instance.database == "testdb"
        assert instance.collection_name == "test_collection"
        assert instance.embedding_model_dims == 128


def test_create_col(azure_mysql_instance):
    """Test collection creation."""
    azure_mysql_instance.create_col(name="new_collection", vector_size=256)

    # Verify that execute was called (table creation)
    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    assert cursor.execute.called


def test_insert(azure_mysql_instance):
    """Test vector insertion."""
    vectors = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
    payloads = [{"text": "test1"}, {"text": "test2"}]
    ids = ["id1", "id2"]

    azure_mysql_instance.insert(vectors=vectors, payloads=payloads, ids=ids)

    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    assert cursor.executemany.called


def test_search(azure_mysql_instance):
    """Test vector search."""
    # Mock the database response
    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    cursor.fetchall = Mock(return_value=[
        {
            'id': 'id1',
            'vector': json.dumps([0.1, 0.2, 0.3]),
            'payload': json.dumps({"text": "test1"})
        },
        {
            'id': 'id2',
            'vector': json.dumps([0.4, 0.5, 0.6]),
            'payload': json.dumps({"text": "test2"})
        }
    ])

    query_vector = [0.2, 0.3, 0.4]
    results = azure_mysql_instance.search(query="test", vectors=query_vector, top_k=5)

    assert isinstance(results, list)
    assert cursor.execute.called


def test_delete(azure_mysql_instance):
    """Test vector deletion."""
    azure_mysql_instance.delete(vector_id="test_id")

    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    assert cursor.execute.called


def test_update(azure_mysql_instance):
    """Test vector update."""
    new_vector = [0.7, 0.8, 0.9]
    new_payload = {"text": "updated"}

    azure_mysql_instance.update(vector_id="test_id", vector=new_vector, payload=new_payload)

    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    assert cursor.execute.called


def test_get(azure_mysql_instance):
    """Test retrieving a vector by ID."""
    # Mock the database response
    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    cursor.fetchone = Mock(return_value={
        'id': 'test_id',
        'vector': json.dumps([0.1, 0.2, 0.3]),
        'payload': json.dumps({"text": "test"})
    })

    result = azure_mysql_instance.get(vector_id="test_id")

    assert result is not None
    assert isinstance(result, OutputData)
    assert result.id == "test_id"


def test_list_cols(azure_mysql_instance):
    """Test listing collections."""
    # Mock the database response
    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    cursor.fetchall = Mock(return_value=[
        {"Tables_in_testdb": "collection1"},
        {"Tables_in_testdb": "collection2"}
    ])

    collections = azure_mysql_instance.list_cols()

    assert isinstance(collections, list)
    assert len(collections) == 2


def test_delete_col(azure_mysql_instance):
    """Test collection deletion."""
    azure_mysql_instance.delete_col()

    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    assert cursor.execute.called


def test_col_info(azure_mysql_instance):
    """Test getting collection information."""
    # Mock the database response
    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    cursor.fetchone = Mock(return_value={
        'name': 'test_collection',
        'count': 100,
        'size_mb': 1.5
    })

    info = azure_mysql_instance.col_info()

    assert isinstance(info, dict)
    assert cursor.execute.called


def test_list(azure_mysql_instance):
    """Test listing vectors."""
    # Mock the database response
    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    cursor.fetchall = Mock(return_value=[
        {
            'id': 'id1',
            'vector': json.dumps([0.1, 0.2, 0.3]),
            'payload': json.dumps({"text": "test1"})
        }
    ])

    results = azure_mysql_instance.list(top_k=10)

    assert isinstance(results, list)
    assert len(results) > 0


def test_reset(azure_mysql_instance):
    """Test resetting the collection."""
    azure_mysql_instance.reset()

    conn = azure_mysql_instance.connection_pool.connection()
    cursor = conn.cursor()
    # Should call execute at least twice (drop and create)
    assert cursor.execute.call_count >= 2


@pytest.mark.skipif(True, reason="Requires Azure credentials")
def test_azure_credential_authentication():
    """Test Azure DefaultAzureCredential authentication."""
    with patch('mem0.vector_stores.azure_mysql.DefaultAzureCredential') as mock_cred:
        mock_token = Mock()
        mock_token.token = "test_token"
        mock_cred.return_value.get_token.return_value = mock_token

        instance = AzureMySQL(
            host="test-server.mysql.database.azure.com",
            port=3306,
            user="testuser",
            password=None,
            database="testdb",
            collection_name="test_collection",
            embedding_model_dims=128,
            use_azure_credential=True,
        )

        assert instance.password == "test_token"


def test_output_data_model():
    """Test OutputData model."""
    data = OutputData(
        id="test_id",
        score=0.95,
        payload={"text": "test"}
    )

    assert data.id == "test_id"
    assert data.score == 0.95
    assert data.payload == {"text": "test"}


INJECTION_PAYLOADS = [
    "memories; DROP TABLE users; --",
    "memories` OR 1=1; --",
    "memories:Label {prop: 'val'}) DELETE n; --",
    "valid_name OR 1=1",
    "1_starts_with_digit",
    "has space",
    "",
]

class TestAzureMySQLConfigCollectionNameValidation:
    def test_accepts_valid_identifier(self):
        config = AzureMySQLConfig(
            collection_name="valid_name",
            host="host",
            user="user",
            database="db",
            password="pw",
        )
        assert config.collection_name == "valid_name"

    @pytest.mark.parametrize("payload", INJECTION_PAYLOADS)
    def test_rejects_injection_payload(self, payload):
        with pytest.raises(ValidationError, match="Invalid collection_name"):
            AzureMySQLConfig(
                collection_name=payload,
                host="host",
                user="user",
                database="db",
                password="pw",
            )


class TestAzureMySQLFilterKeySanitization:
    """Non-identifier filter keys are interpolated into the JSON_EXTRACT path,
    so every method that builds a filter clause must drop them before they run."""

    _FILTERS = {"user_id": "u1", "evil') OR '1'='1": "x"}

    @staticmethod
    def _executed_blob(instance):
        cursor = instance.connection_pool.connection().cursor()
        return " ".join(repr(c.args) + repr(c.kwargs) for c in cursor.execute.call_args_list)

    def test_search_drops_invalid_key(self, azure_mysql_instance):
        azure_mysql_instance.search(query="q", vectors=[0.1, 0.2, 0.3], filters=self._FILTERS)
        blob = self._executed_blob(azure_mysql_instance)
        assert "$.user_id" in blob and "OR '1'='1" not in blob

    def test_keyword_search_drops_invalid_key(self, azure_mysql_instance):
        azure_mysql_instance.keyword_search(query="q", filters=self._FILTERS)
        blob = self._executed_blob(azure_mysql_instance)
        assert "$.user_id" in blob and "OR '1'='1" not in blob

    def test_list_drops_invalid_key(self, azure_mysql_instance):
        azure_mysql_instance.list(filters=self._FILTERS)
        blob = self._executed_blob(azure_mysql_instance)
        assert "$.user_id" in blob and "OR '1'='1" not in blob