Merge branch master of github.com:SyntaxArc/ArchiPy

# Conflicts:
#	poetry.lock

Author: HosseinNejatiJavaremi

archipy/adapters/base/sqlalchemy/adapters.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, override
+from typing import Any, override, TypeVar
 from uuid import UUID
 
 from sqlalchemy import Delete, Executable, Result, ScalarResult, Update, func, select
@@ -30,6 +30,9 @@
 from archipy.models.types.base_types import FilterOperationType
 from archipy.models.types.sort_order_type import SortOrderType
 
+# Generic type variable for BaseEntity subclasses
+T = TypeVar('T', bound=BaseEntity)
+
 
 class SQLAlchemyExceptionHandlerMixin:
     """Mixin providing centralized exception handling for SQLAlchemy operations.
@@ -284,14 +287,14 @@ def get_session(self) -> Session:
         return self.session_manager.get_session()
 
     @override
-    def create(self, entity: BaseEntity) -> BaseEntity | None:
+    def create(self, entity: T) -> T | None:
         """Create a new entity in the database.
 
         Args:
             entity: The entity to create.
 
         Returns:
-            The created entity with updated attributes.
+            The created entity with updated attributes, preserving the original type.
 
         Raises:
             InvalidEntityTypeError: If the entity type is not a valid SQLAlchemy model.
@@ -318,14 +321,14 @@ def create(self, entity: BaseEntity) -> BaseEntity | None:
             return entity
 
     @override
-    def bulk_create(self, entities: list[BaseEntity]) -> list[BaseEntity] | None:
+    def bulk_create(self, entities: list[T]) -> list[T] | None:
         """Creates multiple entities in a single database operation.
 
         Args:
             entities: List of entities to create.
 
         Returns:
-            List of created entities with updated attributes.
+            List of created entities with updated attributes, preserving original types.
 
         Raises:
             InvalidEntityTypeError: If any entity is not a valid SQLAlchemy model.
@@ -352,7 +355,7 @@ def bulk_create(self, entities: list[BaseEntity]) -> list[BaseEntity] | None:
             return entities
 
     @override
-    def get_by_uuid(self, entity_type: type, entity_uuid: UUID) -> BaseEntity | None:
+    def get_by_uuid(self, entity_type: type[T], entity_uuid: UUID) -> T | None:
         """Retrieve an entity by its UUID.
 
         Args:
@@ -593,14 +596,14 @@ def get_session(self) -> AsyncSession:
         return self.session_manager.get_session()
 
     @override
-    async def create(self, entity: BaseEntity) -> BaseEntity | None:
+    async def create(self, entity: T) -> T | None:
         """Create a new entity in the database.
 
         Args:
             entity: The entity to create.
 
         Returns:
-            The created entity with updated attributes.
+            The created entity with updated attributes, preserving the original type.
 
         Raises:
             InvalidEntityTypeError: If the entity type is not a valid SQLAlchemy model.
@@ -627,14 +630,14 @@ async def create(self, entity: BaseEntity) -> BaseEntity | None:
             return entity
 
     @override
-    async def bulk_create(self, entities: list[BaseEntity]) -> list[BaseEntity] | None:
+    async def bulk_create(self, entities: list[T]) -> list[T] | None:
         """Creates multiple entities in a single database operation.
 
         Args:
             entities: List of entities to create.
 
         Returns:
-            List of created entities with updated attributes.
+            List of created entities with updated attributes, preserving original types.
 
         Raises:
             InvalidEntityTypeError: If any entity is not a valid SQLAlchemy model.
@@ -661,7 +664,7 @@ async def bulk_create(self, entities: list[BaseEntity]) -> list[BaseEntity] | No
             return entities
 
     @override
-    async def get_by_uuid(self, entity_type: type, entity_uuid: UUID) -> Any | None:
+    async def get_by_uuid(self, entity_type: type[T], entity_uuid: UUID) -> T | None:
         """Retrieve an entity by its UUID.
 
         Args:

features/environment.py

@@ -65,16 +65,6 @@ def before_scenario(context: Context, scenario: Scenario):
     except Exception as e:
         logger.exception(f"Error setting test config: {e}")
 
-    # Set up async test environment if needed
-    if "async" in scenario.name.lower() or any("async" in tag.lower() for tag in scenario.tags):
-        logger.info("Setting up async test environment")
-        try:
-            # Create a new event loop for this scenario
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            scenario_context.store("_async_test_loop", loop)
-        except Exception as e:
-            logger.exception(f"Error setting up async environment: {e}")
 
 
 def after_scenario(context: Context, scenario: Scenario):

features/scenario_context.py

@@ -43,30 +43,14 @@ def cleanup(self):
         # Clean up async adapter
         if self.async_adapter:
             try:
-                if hasattr(self.async_adapter, "session_manager") and hasattr(
-                    self.async_adapter.session_manager,
-                    "engine",
-                ):
-                    # For async, we need to get a loop and run the coroutine
-                    loop = asyncio.get_event_loop()
-                    if not loop.is_closed():
-                        try:
-                            # Run the session removal coroutine
-                            loop.run_until_complete(self.async_adapter.session_manager.remove_session())
-
-                            # Run the engine disposal coroutine
-                            loop.run_until_complete(self.async_adapter.session_manager.engine.dispose())
-                        except Exception as e:
-                            print(f"Error in async cleanup: {e}")
-                    else:
-                        # If the loop is closed, create a new one temporarily
-                        temp_loop = asyncio.new_event_loop()
-                        try:
-                            asyncio.set_event_loop(temp_loop)
-                            temp_loop.run_until_complete(self.async_adapter.session_manager.remove_session())
-                            temp_loop.run_until_complete(self.async_adapter.session_manager.engine.dispose())
-                        finally:
-                            temp_loop.close()
+                # Try to run async cleanup if we're in an async context
+                try:
+                    loop = asyncio.get_running_loop()
+                    # If we have a running loop, create a task
+                    asyncio.create_task(self.async_cleanup())
+                except RuntimeError:
+                    # No running loop, run in new loop
+                    asyncio.run(self.async_cleanup())
             except Exception as e:
                 print(f"Error in async cleanup: {e}")
 
@@ -80,3 +64,16 @@ def cleanup(self):
                 os.remove(self.db_file)
             except Exception as e:
                 print(f"Error removing database file: {e}")
+
+    async def async_cleanup(self):
+        """Clean up async resources associated with this scenario."""
+        if self.async_adapter:
+            try:
+                if hasattr(self.async_adapter, "session_manager") and hasattr(
+                        self.async_adapter.session_manager, "engine"
+                ):
+                    # Clean up async sessions and engine
+                    await self.async_adapter.session_manager.remove_session()
+                    await self.async_adapter.session_manager.engine.dispose()
+            except Exception as e:
+                print(f"Error in async cleanup: {e}")

features/steps/atomic_transaction_steps.py

@@ -3,7 +3,7 @@
 This module contains step definitions for both synchronous and asynchronous
 atomic transaction scenarios.
 """
-
+import asyncio
 import logging
 import os
 import tempfile
@@ -16,12 +16,10 @@
 from features.test_entity import RelatedTestEntity, TestAdminEntity, TestEntity, TestManagerEntity
 from features.test_entity_factory import TestEntityFactory
 from features.test_helpers import (
-    SafeAsyncContextManager,
     async_schema_setup,
     get_adapter,
     get_async_adapter,
     get_current_scenario_context,
-    safe_run_async,
 )
 from sqlalchemy import select
 
@@ -127,8 +125,7 @@ def step_given_database_initialized(context):
 
             # Create schema with async adapter
             logger.info("Creating database schema with async adapter")
-            with SafeAsyncContextManager(context) as ctx:
-                ctx.run(async_schema_setup(async_adapter))
+            asyncio.run(async_schema_setup(async_adapter))
 
             logger.info("Async adapter and schema setup completed")
         except Exception as e:
@@ -965,7 +962,6 @@ def verify_consistency():
 
 
 @when("a new entity is created in an async atomic transaction")
-@safe_run_async
 async def step_when_entity_created_in_async_atomic(context):
     """Create a new entity within an async atomic transaction."""
     logger = getattr(context, "logger", logging.getLogger("behave.steps"))
@@ -993,7 +989,6 @@ async def create_entity_async_atomic():
 
 
 @then("the async entity should be retrievable")
-@safe_run_async
 async def step_then_async_entity_should_be_retrievable(context):
     """Verify the entity exists after async atomic transaction."""
     logger = getattr(context, "logger", logging.getLogger("behave.steps"))
@@ -1016,7 +1011,6 @@ async def get_entity():
 
 
 @when("a new async entity creation fails within an atomic transaction")
-@safe_run_async
 async def step_when_async_entity_creation_fails(context):
     """Attempt to create an async entity with a failure that causes rollback."""
     logger = getattr(context, "logger", logging.getLogger("behave.steps"))
@@ -1053,7 +1047,6 @@ async def create_entity_with_failure():
 
 
 @then("no async entity should exist in the database")
-@safe_run_async
 async def step_then_no_async_entity_should_exist(context):
     """Verify the entity doesn't exist after failed async atomic transaction."""
     logger = getattr(context, "logger", logging.getLogger("behave.steps"))
@@ -1075,7 +1068,6 @@ async def check_entity_absence():
 
 
 @then("the async database session should remain usable")
-@safe_run_async
 async def step_then_async_session_should_remain_usable(context):
     """Verify the async session is still usable after a failed transaction."""
     logger = getattr(context, "logger", logging.getLogger("behave.steps"))
@@ -1107,7 +1099,6 @@ async def verify_session_usable():
 
 
 @when("multiple entities are created in an async atomic transaction")
-@safe_run_async
 async def step_when_multiple_async_entities_created(context):
     """Create multiple entities in a single async atomic transaction."""
     logger = getattr(context, "logger", logging.getLogger("behave.steps"))
@@ -1141,7 +1132,6 @@ async def create_multiple_entities():
 
 
 @then("all async entities should be retrievable")
-@safe_run_async
 async def step_then_all_async_entities_retrievable(context):
     """Verify all entities exist after async atomic transaction."""
     logger = getattr(context, "logger", logging.getLogger("behave.steps"))
@@ -1177,7 +1167,6 @@ async def verify_entities():
 
 
 @when("complex async operations are performed in a transaction")
-@safe_run_async
 async def step_when_complex_async_operations(context):
     """Demonstrate more complex async operations with proper session management."""
     logger = getattr(context, "logger", logging.getLogger("behave.steps"))
@@ -1227,7 +1216,6 @@ async def create_entity_with_relations():
 
 
 @then("all related entities should be accessible")
-@safe_run_async
 async def step_then_related_entities_accessible(context):
     """Verify that related entities can be accessed through relationships."""
     logger = getattr(context, "logger", logging.getLogger("behave.steps"))