chore: Add tests

Also return data from an insert with a returning statement.

Author: jonyscathe

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+  "python.testing.pytestArgs": ["tests"],
+  "python.testing.unittestEnabled": false,
+  "python.testing.pytestEnabled": true
+}

src/mock_alchemy/comparison.py

@@ -10,8 +10,8 @@
 from typing import Optional
 from unittest import mock
 
-import sqlalchemy
 from packaging import version
+from sqlalchemy import __version__ as sqlalchemy_version
 from sqlalchemy import delete
 from sqlalchemy import func
 from sqlalchemy import insert
@@ -35,7 +35,7 @@
     ALCHEMY_FUNC_TYPE,
     ALCHEMY_LABEL_TYPE,
 )
-if version.parse(sqlalchemy.__version__) >= version.parse("1.4.0"):
+if version.parse(sqlalchemy_version) >= version.parse("1.4.0"):
     ALCHEMY_SELECT_TYPE = type(select(column("")))
     ALCHEMY_UPDATE_TYPE = type(update(table("")))
     ALCHEMY_DELETE_TYPE = type(delete(table("")))

src/mock_alchemy/mocking.py

@@ -17,8 +17,9 @@
 from typing import overload
 from unittest import mock
 
+from packaging import version
+from sqlalchemy import __version__ as sqlalchemy_version
 from sqlalchemy import select
-from sqlalchemy.exc import ArgumentError
 from sqlalchemy.orm.exc import MultipleResultsFound
 from sqlalchemy.orm.exc import NoResultFound
 from sqlalchemy.sql.dml import Delete
@@ -459,7 +460,6 @@ class UnifiedAlchemyMagicMock(AlchemyMagicMock):
     unify: Dict[str, Optional[UnorderedCall]] = {
         "add_columns": None,
         "distinct": None,
-        "execute": None,
         "filter": UnorderedCall,
         "filter_by": UnorderedCall,
         "group_by": None,
@@ -473,9 +473,9 @@ class UnifiedAlchemyMagicMock(AlchemyMagicMock):
         "where": None,
     }
 
-    mutate: Set[str] = {"add", "add_all", "delete", "execute"}
+    mutate: Set[str] = {"add", "add_all", "delete"}
 
-    mutate_and_unify: Set[str] = {"execute"}
+    execute_statement: Set[str] = {"execute"}
 
     @overload
     def __init__(
@@ -498,7 +498,12 @@ def __init__(self, *args, **kwargs) -> None:
         """Creates an UnifiedAlchemyMagicMock to mock a SQLAlchemy session."""
         kwargs["_mock_default"] = kwargs.pop("default", [])
         kwargs["_mock_data"] = kwargs.pop("data", None)
-        kwargs.update({k: AlchemyMagicMock(side_effect=partial(self._get_data, _mock_name=k)) for k in self.boundary})
+        kwargs.update(
+            {
+                k: AlchemyMagicMock(side_effect=partial(self._get_data, _mock_name=k))
+                for k in self.boundary
+            }
+        )
 
         kwargs.update(
             {
@@ -524,9 +529,9 @@ def __init__(self, *args, **kwargs) -> None:
             {
                 k: AlchemyMagicMock(
                     return_value=self,
-                    side_effect=partial(self._mutate_data, _mock_name=k),
+                    side_effect=partial(self._execute_statement, _mock_name=k),
                 )
-                for k in self.mutate_and_unify
+                for k in self.execute_statement
             }
         )
 
@@ -604,13 +609,25 @@ def _get_data(self, *args: Any, **kwargs: Any) -> Any:
         _mock_data = self._mock_data
         if _mock_data is not None:
             previous_calls = [
-                sqlalchemy_call(i, with_name=True, base_call=self.unify.get(i[0]) or Call)
+                sqlalchemy_call(
+                    i, with_name=True, base_call=self.unify.get(i[0]) or Call
+                )
                 for i in self._get_previous_calls(self.mock_calls[:-1])
             ]
             sorted_mock_data = sorted(_mock_data, key=lambda x: len(x[0]), reverse=True)
             if _mock_name == "get":
-                query_call = [c for c in previous_calls if c[0] in ["query", "execute"]][0]
-                results = list(chain(*[result for calls, result in sorted_mock_data if query_call in calls]))
+                query_call = [
+                    c for c in previous_calls if c[0] in ["query", "execute"]
+                ][0]
+                results = list(
+                    chain(
+                        *[
+                            result
+                            for calls, result in sorted_mock_data
+                            if query_call in calls
+                        ]
+                    )
+                )
                 return self.boundary[_mock_name](results, *args, **kwargs)
 
             else:
@@ -636,13 +653,15 @@ def _mutate_data(self, *args: Any, **kwargs: Any) -> Optional[int]:
             to_add = args[0]
             query_call = mock.call.query(type(to_add))
 
-            mocked_data = next(iter(filter(lambda i: i[0] == [query_call], _mock_data)), None)
+            mocked_data = next(
+                iter(filter(lambda i: i[0] == [query_call], _mock_data)), None
+            )
             if mocked_data:
                 mocked_data[1].append(to_add)
             else:
                 _mock_data.append(([query_call], [to_add]))
 
-            try:
+            if version.parse(sqlalchemy_version) >= version.parse("1.4.0"):
                 execute_call = mock.call.execute(select(type(to_add)))
 
                 execute_mocked_data = next(
@@ -658,29 +677,25 @@ def _mutate_data(self, *args: Any, **kwargs: Any) -> Optional[int]:
                     execute_mocked_data[1].append(to_add)
                 else:
                     _mock_data.append(([execute_call], [to_add]))
-            except (TypeError, ArgumentError):
-                # TypeError indicates an old version of sqlalcemy that does not support
-                # executing select(table) statements.
-                # ArgumentError indicates a mocked table that cannot be selected, so
-                # cannot be mocked this way.
-                pass
-
         elif _mock_name == "add_all":
             to_add = args[0]
             _kwargs = kwargs.copy()
             _kwargs["_mock_name"] = "add"
 
             for i in to_add:
                 self._mutate_data(i, *args[1:], **_kwargs)
-        elif _mock_name == "delete":
+        # delete case
+        else:
             _kwargs = kwargs.copy()
             # pretend like all is being called to get data
             _kwargs["_mock_name"] = "all"
             _mock_name = _kwargs.pop("_mock_name")
             _mock_data = self._mock_data
             num_deleted = 0
             previous_calls = [
-                sqlalchemy_call(i, with_name=True, base_call=self.unify.get(i[0]) or Call)
+                sqlalchemy_call(
+                    i, with_name=True, base_call=self.unify.get(i[0]) or Call
+                )
                 for i in self._get_previous_calls(self.mock_calls[:-1])
             ]
             sorted_mock_data = sorted(_mock_data, key=lambda x: len(x[0]), reverse=True)
@@ -703,100 +718,134 @@ def _mutate_data(self, *args: Any, **kwargs: Any) -> Optional[int]:
                     temp_mock_data.append((calls, result))
             self._mock_data = temp_mock_data
             return num_deleted
-        # execute case
+
+    def _execute_insert(
+        self, execute_statement: Insert, *args: Any, **kwargs: Any
+    ) -> Any:
+        """Insert data from execute statement."""
+        _kwargs = kwargs.copy()
+        execute_statement = args[0]
+        _kwargs["_mock_name"] = "add"
+        table_type = execute_statement.entity_description["type"]
+        # Values should either be a list of dictionaries as arg[1] or a list of
+        # dictionaries as values.
+        if len(args) > 1:
+            for i in args[1]:
+                self._mutate_data(table_type(**i), **_kwargs)
         else:
-            _kwargs = kwargs.copy()
-            # Need to check if the execute was an insert, update or delete. Ignore any other types
-            execute_statement = args[0]
-
-            if isinstance(execute_statement, Insert):
-                # Add insert data
-                _kwargs["_mock_name"] = "add"
-                table_type = execute_statement.entity_description["type"]
-                # Values should either be a list of dictionaries ar arg[1] or a list of dictionaries as values.
-                if len(args) > 1:
-                    for i in args[1]:
-                        self._mutate_data(table_type(**i), **_kwargs)
-                else:
-                    # Values will be stored within _multi_values list
-                    values = execute_statement._multi_values[0]
-                    for i in values:
-                        self._mutate_data(table_type(**{k.name: v for k, v in i.items()}), **_kwargs)
-                # Only unify if the insert statement is returning
-                if execute_statement._returning:
-                    return self._unify(self, *args, **kwargs)
-                else:
-                    # insert a boundary so that this is no longer part of a unified call.
-                    self.all()
-            elif isinstance(execute_statement, Delete):
-                # Create equivalent select statement as an Expression Matcher
-                select_statement = [
-                    ExpressionMatcher(
-                        mock.call.execute(select(execute_statement.table).where(execute_statement.whereclause))
-                    )
-                ]
-                _mock_data = self._mock_data
-                sorted_mock_data = sorted(_mock_data, key=lambda x: len(x[0]), reverse=True)
-                temp_mock_data = list()
-                found_query = False
-                num_deleted = 0
-                for calls, result in sorted_mock_data:
-                    calls = [
-                        sqlalchemy_call(
-                            i,
-                            with_name=True,
-                            base_call=self.unify.get(i[0]) or Call,
+            # Values will be stored within _multi_values list
+            values = execute_statement._multi_values[0]
+            for i in values:
+                self._mutate_data(
+                    table_type(**{k.name: v for k, v in i.items()}), **_kwargs
+                )
+        # insert a boundary so that this is no longer part of a unified call.
+        self.all()
+        # Start a new unify if the insert statement is returning
+        if execute_statement._returning:
+            return self.execute(select(execute_statement._returning[0]))
+        return None
+
+    def _execute_delete(self, execute_statement: Delete, *args: Any) -> mock.Mock:
+        """Delete data according to execute statement."""
+        execute_statement = args[0]
+        # Create equivalent select statement as an Expression Matcher
+        select_statement = (
+            [
+                ExpressionMatcher(
+                    mock.call.execute(
+                        select(execute_statement.table).where(
+                            execute_statement.whereclause
                         )
-                        for i in calls
-                    ]
-                    if all(c in select_statement for c in calls) and not found_query:
-                        num_deleted = len(result)
-                        temp_mock_data.append((calls, []))
-                        found_query = True
-                    else:
-                        temp_mock_data.append((calls, result))
-                self._mock_data = temp_mock_data
-                delete_result = mock.Mock()
-                delete_result.rowcount = num_deleted
-                # insert a boundary so that this is no longer part of a unified call.
-                self.all()
-                return delete_result
-            elif isinstance(execute_statement, Update):
-                # Create equivalent select statement as an Expression Matcher
-                select_statement = [
-                    ExpressionMatcher(
-                        mock.call.execute(select(execute_statement.table).where(execute_statement.whereclause))
                     )
-                ]
-                _mock_data = self._mock_data
-                sorted_mock_data = sorted(_mock_data, key=lambda x: len(x[0]), reverse=True)
-                temp_mock_data = list()
-                found_query = False
-                num_updated = 0
-                for calls, result in sorted_mock_data:
-                    calls = [
-                        sqlalchemy_call(
-                            i,
-                            with_name=True,
-                            base_call=self.unify.get(i[0]) or Call,
+                )
+            ]
+            if execute_statement.whereclause is not None
+            else [ExpressionMatcher(mock.call.execute(select(execute_statement.table)))]
+        )
+        _mock_data = self._mock_data = self._mock_data or []
+        sorted_mock_data = sorted(_mock_data, key=lambda x: len(x[0]), reverse=True)
+        temp_mock_data = list()
+        found_query = False
+        num_deleted = 0
+        for calls, result in sorted_mock_data:
+            calls = [
+                sqlalchemy_call(
+                    i,
+                    with_name=True,
+                    base_call=self.unify.get(i[0]) or Call,
+                )
+                for i in calls
+            ]
+            if all(c in select_statement for c in calls) and not found_query:
+                num_deleted = len(result)
+                temp_mock_data.append((calls, []))
+                found_query = True
+            else:
+                temp_mock_data.append((calls, result))
+        self._mock_data = temp_mock_data
+        delete_result = mock.Mock()
+        delete_result.rowcount = num_deleted
+        # insert a boundary so that this is no longer part of a unified call.
+        self.all()
+        return delete_result
+
+    def _execute_update(self, execute_statement: Update) -> mock.Mock:
+        """Update data according to execute statement."""
+        # Create equivalent select statement as an Expression Matcher
+        select_statement = (
+            [
+                ExpressionMatcher(
+                    mock.call.execute(
+                        select(execute_statement.table).where(
+                            execute_statement.whereclause
                         )
-                        for i in calls
-                    ]
-                    if all(c in select_statement for c in calls) and not found_query:
-                        num_updated = len(result)
-                        for r in result:
-                            for k, v in execute_statement._values.items():
-                                setattr(r, k.name, v.value)
-                        temp_mock_data.append((calls, result))
-                        found_query = True
-                    else:
-                        temp_mock_data.append((calls, result))
-                self._mock_data = temp_mock_data
-                update_result = mock.Mock()
-                update_result.rowcount = num_updated
-                # insert a boundary so that this is no longer part of a unified call.
-                self.all()
-                return update_result
+                    )
+                )
+            ]
+            if execute_statement.whereclause is not None
+            else [ExpressionMatcher(mock.call.execute(select(execute_statement.table)))]
+        )
+        _mock_data = self._mock_data = self._mock_data or []
+        sorted_mock_data = sorted(_mock_data, key=lambda x: len(x[0]), reverse=True)
+        temp_mock_data = list()
+        found_query = False
+        num_updated = 0
+        for calls, result in sorted_mock_data:
+            calls = [
+                sqlalchemy_call(
+                    i,
+                    with_name=True,
+                    base_call=self.unify.get(i[0]) or Call,
+                )
+                for i in calls
+            ]
+            if all(c in select_statement for c in calls) and not found_query:
+                num_updated = len(result)
+                for r in result:
+                    for k, v in execute_statement._values.items():
+                        setattr(r, k.name, v.value)
+                temp_mock_data.append((calls, result))
+                found_query = True
             else:
-                # assume any other execute types need to unify
-                return self._unify(self, *args, **kwargs)
+                temp_mock_data.append((calls, result))
+        self._mock_data = temp_mock_data
+        update_result = mock.Mock()
+        update_result.rowcount = num_updated
+        # insert a boundary so that this is no longer part of a unified call.
+        self.all()
+        return update_result
+
+    def _execute_statement(self, *args: Any, **kwargs: Any) -> Any:
+        """Depending on statement being executed, update data and/or unify statement."""
+        # Need to check if the execute was an insert, update or delete.
+        execute_statement = args[0]
+        if isinstance(execute_statement, Insert):
+            return self._execute_insert(execute_statement, *args, **kwargs)
+        elif isinstance(execute_statement, Delete):
+            return self._execute_delete(execute_statement, *args)
+        elif isinstance(execute_statement, Update):
+            return self._execute_update(execute_statement)
+        else:
+            # assume any other execute types need to unify
+            return self._unify(self, *args, **kwargs)