cocoindex 0.2.3__cp311-abi3-macosx_10_12_x86_64.whl

cocoindex/tests/test_optional_database.py

@@ -0,0 +1,249 @@
+"""
+Test suite for optional database functionality in CocoIndex.
+
+This module tests that:
+1. cocoindex.init() works without database settings
+2. Transform flows work without database
+3. Database functionality still works when database settings are provided
+4. Operations requiring database properly complain when no database is configured
+"""
+
+import os
+from unittest.mock import patch
+import pytest
+
+import cocoindex
+from cocoindex import op
+from cocoindex.setting import Settings
+
+
+class TestOptionalDatabase:
+    """Test suite for optional database functionality."""
+
+    def setup_method(self) -> None:
+        """Setup method called before each test."""
+        # Stop any existing cocoindex instance
+        try:
+            cocoindex.stop()
+        except:
+            pass
+
+    def teardown_method(self) -> None:
+        """Teardown method called after each test."""
+        # Stop cocoindex instance after each test
+        try:
+            cocoindex.stop()
+        except:
+            pass
+
+    def test_init_without_database(self) -> None:
+        """Test that cocoindex.init() works without database settings."""
+        # Remove database environment variables
+        with patch.dict(os.environ, {}, clear=False):
+            # Remove database env vars if they exist
+            for env_var in [
+                "COCOINDEX_DATABASE_URL",
+                "COCOINDEX_DATABASE_USER",
+                "COCOINDEX_DATABASE_PASSWORD",
+            ]:
+                os.environ.pop(env_var, None)
+
+            # Test initialization without database
+            cocoindex.init()
+
+            # If we get here without exception, the test passes
+            assert True
+
+    def test_transform_flow_without_database(self) -> None:
+        """Test that transform flows work without database."""
+        # Remove database environment variables
+        with patch.dict(os.environ, {}, clear=False):
+            # Remove database env vars if they exist
+            for env_var in [
+                "COCOINDEX_DATABASE_URL",
+                "COCOINDEX_DATABASE_USER",
+                "COCOINDEX_DATABASE_PASSWORD",
+            ]:
+                os.environ.pop(env_var, None)
+
+            # Initialize without database
+            cocoindex.init()
+
+            # Create a simple custom function for testing
+            @op.function()
+            def add_prefix(text: str) -> str:
+                """Add a prefix to text."""
+                return f"processed: {text}"
+
+            @cocoindex.transform_flow()
+            def simple_transform(
+                text: cocoindex.DataSlice[str],
+            ) -> cocoindex.DataSlice[str]:
+                """A simple transform that adds a prefix."""
+                return text.transform(add_prefix)
+
+            # Test the transform flow
+            result = simple_transform.eval("hello world")
+            expected = "processed: hello world"
+
+            assert result == expected
+
+    @pytest.mark.skipif(
+        not os.getenv("COCOINDEX_DATABASE_URL"),
+        reason="Database URL not configured in environment",
+    )
+    def test_init_with_database(self) -> None:
+        """Test that cocoindex.init() works with database settings when available."""
+        # This test only runs if database URL is configured
+        settings = Settings.from_env()
+        assert settings.database is not None
+        assert settings.database.url is not None
+
+        try:
+            cocoindex.init(settings)
+            assert True
+        except Exception as e:
+            assert (
+                "Failed to connect to database" in str(e)
+                or "connection" in str(e).lower()
+            )
+
+    def test_settings_from_env_without_database(self) -> None:
+        """Test that Settings.from_env() correctly handles missing database settings."""
+        with patch.dict(os.environ, {}, clear=False):
+            # Remove database env vars if they exist
+            for env_var in [
+                "COCOINDEX_DATABASE_URL",
+                "COCOINDEX_DATABASE_USER",
+                "COCOINDEX_DATABASE_PASSWORD",
+            ]:
+                os.environ.pop(env_var, None)
+
+            settings = Settings.from_env()
+            assert settings.database is None
+            assert settings.app_namespace == ""
+
+    def test_settings_from_env_with_database(self) -> None:
+        """Test that Settings.from_env() correctly handles database settings when provided."""
+        test_url = "postgresql://test:test@localhost:5432/test"
+        test_user = "testuser"
+        test_password = "testpass"
+
+        with patch.dict(
+            os.environ,
+            {
+                "COCOINDEX_DATABASE_URL": test_url,
+                "COCOINDEX_DATABASE_USER": test_user,
+                "COCOINDEX_DATABASE_PASSWORD": test_password,
+            },
+        ):
+            settings = Settings.from_env()
+            assert settings.database is not None
+            assert settings.database.url == test_url
+            assert settings.database.user == test_user
+            assert settings.database.password == test_password
+
+    def test_settings_from_env_with_partial_database_config(self) -> None:
+        """Test Settings.from_env() with only database URL (no user/password)."""
+        test_url = "postgresql://localhost:5432/test"
+
+        with patch.dict(
+            os.environ,
+            {
+                "COCOINDEX_DATABASE_URL": test_url,
+            },
+            clear=False,
+        ):
+            # Remove user/password env vars if they exist
+            os.environ.pop("COCOINDEX_DATABASE_USER", None)
+            os.environ.pop("COCOINDEX_DATABASE_PASSWORD", None)
+
+            settings = Settings.from_env()
+            assert settings.database is not None
+            assert settings.database.url == test_url
+            assert settings.database.user is None
+            assert settings.database.password is None
+
+    def test_multiple_init_calls(self) -> None:
+        """Test that multiple init calls work correctly."""
+        with patch.dict(os.environ, {}, clear=False):
+            # Remove database env vars if they exist
+            for env_var in [
+                "COCOINDEX_DATABASE_URL",
+                "COCOINDEX_DATABASE_USER",
+                "COCOINDEX_DATABASE_PASSWORD",
+            ]:
+                os.environ.pop(env_var, None)
+
+            # First init
+            cocoindex.init()
+
+            # Stop and init again
+            cocoindex.stop()
+            cocoindex.init()
+
+            # Should work without issues
+            assert True
+
+    def test_app_namespace_setting(self) -> None:
+        """Test that app_namespace setting works correctly."""
+        test_namespace = "test_app"
+
+        with patch.dict(
+            os.environ,
+            {
+                "COCOINDEX_APP_NAMESPACE": test_namespace,
+            },
+            clear=False,
+        ):
+            # Remove database env vars if they exist
+            for env_var in [
+                "COCOINDEX_DATABASE_URL",
+                "COCOINDEX_DATABASE_USER",
+                "COCOINDEX_DATABASE_PASSWORD",
+            ]:
+                os.environ.pop(env_var, None)
+
+            settings = Settings.from_env()
+            assert settings.app_namespace == test_namespace
+            assert settings.database is None
+
+            # Init should work with app namespace but no database
+            cocoindex.init(settings)
+            assert True
+
+
+class TestDatabaseRequiredOperations:
+    """Test suite for operations that require database."""
+
+    def setup_method(self) -> None:
+        """Setup method called before each test."""
+        # Stop any existing cocoindex instance
+        try:
+            cocoindex.stop()
+        except:
+            pass
+
+    def teardown_method(self) -> None:
+        """Teardown method called after each test."""
+        # Stop cocoindex instance after each test
+        try:
+            cocoindex.stop()
+        except:
+            pass
+
+    def test_database_required_error_message(self) -> None:
+        """Test that operations requiring database show proper error messages."""
+        with patch.dict(os.environ, {}, clear=False):
+            # Remove database env vars if they exist
+            for env_var in [
+                "COCOINDEX_DATABASE_URL",
+                "COCOINDEX_DATABASE_USER",
+                "COCOINDEX_DATABASE_PASSWORD",
+            ]:
+                os.environ.pop(env_var, None)
+
+            # Initialize without database
+            cocoindex.init()
+
+            assert True

cocoindex/tests/test_transform_flow.py

@@ -0,0 +1,207 @@
+import typing
+from dataclasses import dataclass
+from typing import Any
+
+import pytest
+
+import cocoindex
+
+
+@dataclass
+class Child:
+    value: int
+
+
+@dataclass
+class Parent:
+    children: list[Child]
+
+
+# Fixture to initialize CocoIndex library
+@pytest.fixture(scope="session", autouse=True)
+def init_cocoindex() -> typing.Generator[None, None, None]:
+    cocoindex.init()
+    yield
+
+
+@cocoindex.op.function()
+def add_suffix(text: str) -> str:
+    """Append ' world' to the input text."""
+    return f"{text} world"
+
+
+@cocoindex.transform_flow()
+def simple_transform(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[str]:
+    """Transform flow that applies add_suffix to input text."""
+    return text.transform(add_suffix)
+
+
+@cocoindex.op.function()
+def extract_value(value: int) -> int:
+    """Extracts the value."""
+    return value
+
+
+@cocoindex.transform_flow()
+def for_each_transform(
+    data: cocoindex.DataSlice[Parent],
+) -> cocoindex.DataSlice[Any]:
+    """Transform flow that processes child rows to extract values."""
+    with data["children"].row() as child:
+        child["new_field"] = child["value"].transform(extract_value)
+    return data
+
+
+def test_simple_transform_flow() -> None:
+    """Test the simple transform flow."""
+    input_text = "hello"
+    result = simple_transform.eval(input_text)
+    assert result == "hello world", f"Expected 'hello world', got {result}"
+
+    result = simple_transform.eval("")
+    assert result == " world", f"Expected ' world', got {result}"
+
+
+@pytest.mark.asyncio
+async def test_simple_transform_flow_async() -> None:
+    """Test the simple transform flow asynchronously."""
+    input_text = "async"
+    result = await simple_transform.eval_async(input_text)
+    assert result == "async world", f"Expected 'async world', got {result}"
+
+
+def test_for_each_transform_flow() -> None:
+    """Test the complex transform flow with child rows."""
+    input_data = Parent(children=[Child(1), Child(2), Child(3)])
+    result = for_each_transform.eval(input_data)
+    expected = {
+        "children": [
+            {"value": 1, "new_field": 1},
+            {"value": 2, "new_field": 2},
+            {"value": 3, "new_field": 3},
+        ]
+    }
+    assert result == expected, f"Expected {expected}, got {result}"
+
+    input_data = Parent(children=[])
+    result = for_each_transform.eval(input_data)
+    assert result == {"children": []}, f"Expected {{'children': []}}, got {result}"
+
+
+@pytest.mark.asyncio
+async def test_for_each_transform_flow_async() -> None:
+    """Test the complex transform flow asynchronously."""
+    input_data = Parent(children=[Child(4), Child(5)])
+    result = await for_each_transform.eval_async(input_data)
+    expected = {
+        "children": [
+            {"value": 4, "new_field": 4},
+            {"value": 5, "new_field": 5},
+        ]
+    }
+
+    assert result == expected, f"Expected {expected}, got {result}"
+
+
+def test_none_arg_yield_none_result() -> None:
+    """Test that None arguments yield None results."""
+
+    @cocoindex.op.function()
+    def custom_fn(
+        required_arg: int,
+        optional_arg: int | None,
+        required_kwarg: int,
+        optional_kwarg: int | None,
+    ) -> int:
+        return (
+            required_arg + (optional_arg or 0) + required_kwarg + (optional_kwarg or 0)
+        )
+
+    @cocoindex.transform_flow()
+    def transform_flow(
+        required_arg: cocoindex.DataSlice[int | None],
+        optional_arg: cocoindex.DataSlice[int | None],
+        required_kwarg: cocoindex.DataSlice[int | None],
+        optional_kwarg: cocoindex.DataSlice[int | None],
+    ) -> cocoindex.DataSlice[int | None]:
+        return required_arg.transform(
+            custom_fn,
+            optional_arg,
+            required_kwarg=required_kwarg,
+            optional_kwarg=optional_kwarg,
+        )
+
+    result = transform_flow.eval(1, 2, 4, 8)
+    assert result == 15, f"Expected 15, got {result}"
+
+    result = transform_flow.eval(1, None, 4, None)
+    assert result == 5, f"Expected 5, got {result}"
+
+    result = transform_flow.eval(None, 2, 4, 8)
+    assert result is None, f"Expected None, got {result}"
+
+    result = transform_flow.eval(1, 2, None, None)
+    assert result is None, f"Expected None, got {result}"
+
+
+# Test GPU function behavior.
+# They're not really executed on GPU, but we want to make sure they're scheduled on subprocesses correctly.
+
+
+@cocoindex.op.function(gpu=True)
+def gpu_append_world(text: str) -> str:
+    """Append ' world' to the input text."""
+    return f"{text} world"
+
+
+class GpuAppendSuffix(cocoindex.op.FunctionSpec):
+    suffix: str
+
+
+@cocoindex.op.executor_class(gpu=True)
+class GpuAppendSuffixExecutor:
+    spec: GpuAppendSuffix
+
+    def __call__(self, text: str) -> str:
+        return f"{text}{self.spec.suffix}"
+
+
+class GpuAppendSuffixWithAnalyzePrepare(cocoindex.op.FunctionSpec):
+    suffix: str
+
+
+@cocoindex.op.executor_class(gpu=True)
+class GpuAppendSuffixWithAnalyzePrepareExecutor:
+    spec: GpuAppendSuffixWithAnalyzePrepare
+    suffix: str
+
+    def analyze(self) -> Any:
+        return str
+
+    def prepare(self) -> None:
+        self.suffix = self.spec.suffix
+
+    def __call__(self, text: str) -> str:
+        return f"{text}{self.suffix}"
+
+
+def test_gpu_function() -> None:
+    @cocoindex.transform_flow()
+    def transform_flow(text: cocoindex.DataSlice[str]) -> cocoindex.DataSlice[str]:
+        return text.transform(gpu_append_world).transform(GpuAppendSuffix(suffix="!"))
+
+    result = transform_flow.eval("Hello")
+    expected = "Hello world!"
+    assert result == expected, f"Expected {expected}, got {result}"
+
+    @cocoindex.transform_flow()
+    def transform_flow_with_analyze_prepare(
+        text: cocoindex.DataSlice[str],
+    ) -> cocoindex.DataSlice[str]:
+        return text.transform(gpu_append_world).transform(
+            GpuAppendSuffixWithAnalyzePrepare(suffix="!!")
+        )
+
+    result = transform_flow_with_analyze_prepare.eval("Hello")
+    expected = "Hello world!!"
+    assert result == expected, f"Expected {expected}, got {result}"