Repository: mem0ai/mem0
File: tests/vectordb/test_zilliz_db.py (Python)

Code before:
# ruff: noqa: E501

import os
from dotenv import load_dotenv

import unittest
from unittest import mock
from unittest.mock import patch

from embedchain import App
from embedchain.config import ZillizDBConfig
from embedchain.vectordb.zilliz import ZillizVectorDB

load_dotenv()


# to run tests, provide the URI and TOKEN in .env file
class TestZillizVectorDBClient(unittest.TestCase):
    @mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
    def test_init_with_uri_and_token(self):
        """
        Test if the `ZillizVectorDB` instance is initialized with the correct uri and token values.
        """
        # Create a ZillizDBConfig instance with mocked values
        expected_uri = "mocked_uri"
        expected_token = "mocked_token"
        db_config = ZillizDBConfig()

        # Assert that the values in the ZillizVectorDB instance match the mocked values
        self.assertEqual(db_config.uri, expected_uri)
        self.assertEqual(db_config.token, expected_token)

    @mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
    def test_init_without_uri(self):
        # Make sure it's not loaded from env
        try:
            del os.environ["ZILLIZ_CLOUD_URI"]
        except KeyError:
            pass
        # Test if an exception is raised when ZILLIZ_CLOUD_URI is missing
        with self.assertRaises(AttributeError):
            ZillizDBConfig()
            

Review comment (security):
Can you please use pytest instead of using `unittest`? We are moving towards using pytest alone. 

Moreover, seems like these tests only test the `ZillizDBConfig` and not the db class itself? Can you please add tests for that as well? You may want to mock the actual zilliz client so that we don't need the credentials for testing. See tests for other dbs for reference. 

Code after:
# ruff: noqa: E501

import os

import pytest
from unittest import mock
from unittest.mock import patch, Mock

from embedchain.config import ZillizDBConfig
from embedchain.vectordb.zilliz import ZillizVectorDB

# to run tests, provide the URI and TOKEN in .env file
class TestZillizVectorDBConfig:
    @mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
    def test_init_with_uri_and_token(self):
        """
        Test if the `ZillizVectorDBConfig` instance is initialized with the correct uri and token values.
        """
        # Create a ZillizDBConfig instance with mocked values
        expected_uri = "mocked_uri"
        expected_token = "mocked_token"
        db_config = ZillizDBConfig()

        # Assert that the values in the ZillizVectorDB instance match the mocked values
        assert db_config.uri == expected_uri
        assert db_config.token == expected_token

    @mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
    def test_init_without_uri(self):
        """
        Test if the `ZillizVectorDBConfig` instance throws an error when no URI found.
        """
        try:
            del os.environ["ZILLIZ_CLOUD_URI"]
        except KeyError:
            pass

        with pytest.raises(AttributeError):
            ZillizDBConfig()

    @mock.patch.dict(os.environ, {"ZILLIZ_CLOUD_URI": "mocked_uri", "ZILLIZ_CLOUD_TOKEN": "mocked_token"})
    def test_init_without_token(self):
        """