change error

dagster-io · Dec 19, 2023 · 9569d69 · 9569d69
1 parent 0f49350
commit 9569d69
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/python_modules/dagster/dagster/_core/storage/db_io_manager.py b/python_modules/dagster/dagster/_core/storage/db_io_manager.py
@@ -26,7 +26,7 @@
     TimeWindow,
     TimeWindowPartitionsDefinition,
 )
-from dagster._core.errors import DagsterExecutionHandleOutputError
+from dagster._core.errors import DagsterInvariantViolationError
 from dagster._core.execution.context.input import InputContext
 from dagster._core.execution.context.output import OutputContext
 from dagster._core.storage.io_manager import IOManager
@@ -125,9 +125,8 @@ def handle_output(self, context: OutputContext, obj: object) -> None:
         # Nones cannot be stored by DB I/O managers. If the output type is set to None/Nothing,
         # the handle_output will not be called
         if obj is None:
-            raise DagsterExecutionHandleOutputError(
-                "Unexpected 'None' output value. If a 'None' value is intentional, set the"
-                " output type to None.",
+            raise DagsterInvariantViolationError(
+                "Unexpected 'None' output value. If a 'None' value is intentional, set the output type to None.",
             )
 
         table_slice = self._get_table_slice(context, context)

diff --git a/python_modules/dagster/dagster_tests/storage_tests/test_db_io_manager.py b/python_modules/dagster/dagster_tests/storage_tests/test_db_io_manager.py
@@ -5,6 +5,7 @@
 from dagster._check import CheckError
 from dagster._core.definitions.partition import StaticPartitionsDefinition
 from dagster._core.definitions.time_window_partitions import DailyPartitionsDefinition, TimeWindow
+from dagster._core.errors import DagsterInvariantViolationError
 from dagster._core.storage.db_io_manager import (
     DbClient,
     DbIOManager,
@@ -483,15 +484,17 @@ def test_handle_none_output():
     handler = IntHandler()
     db_client = MagicMock(spec=DbClient, get_select_statement=MagicMock(return_value=""))
     manager = build_db_io_manager(type_handlers=[handler], db_client=db_client)
+
     asset_key = AssetKey(["schema1", "table1"])
     output_context = build_output_context(
         asset_key=asset_key,
         resource_config=resource_config,
         dagster_type=resolve_dagster_type(type(None)),
+        name="result",
     )
-    manager.handle_output(output_context, None)
 
-    assert len(handler.handle_output_calls) == 0
+    with pytest.raises(DagsterInvariantViolationError):
+        manager.handle_output(output_context, None)
 
 
 def test_non_supported_type():