From 3c36151402661b8dae48011a597fd4cf51ac8dd8 Mon Sep 17 00:00:00 2001
From: Waket Zheng <waketzheng@gmail.com>
Date: Sat, 15 Jun 2024 01:42:04 +0800
Subject: [PATCH] Fix breaking change of `DoesNotExist` (#1650)

* Support str argument for DoesNotExist

* make style
---
 tests/test_queryset.py | 16 ++++++++++++++++
 tortoise/exceptions.py | 37 ++++++++++++++++++++++---------------
 2 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/tests/test_queryset.py b/tests/test_queryset.py
index d7d49faf8..32c9af2fb 100644
--- a/tests/test_queryset.py
+++ b/tests/test_queryset.py
@@ -1,3 +1,5 @@
+from typing import Type
+
 from tests.testmodels import (
     Event,
     IntFields,
@@ -16,6 +18,7 @@
     IntegrityError,
     MultipleObjectsReturned,
     ParamsError,
+    NotExistOrMultiple,
 )
 from tortoise.expressions import F, RawSQL, Subquery
 
@@ -673,3 +676,16 @@ async def test_annotation_field_priorior_to_model_field(self):
         t1 = await Tournament.create(name="1")
         ret = await Tournament.filter(pk=t1.pk).annotate(id=RawSQL("id + 1")).values("id")
         self.assertEqual(ret, [{"id": t1.pk + 1}])
+
+
+class TestNotExist(test.TestCase):
+    exp_cls: Type[NotExistOrMultiple] = DoesNotExist
+
+    @test.requireCapability(dialect="sqlite")
+    def test_does_not_exist(self):
+        assert str(self.exp_cls("old format")) == "old format"
+        assert str(self.exp_cls(Tournament)) == self.exp_cls.TEMPLATE.format(Tournament.__name__)
+
+
+class TestMultiple(TestNotExist):
+    exp_cls = MultipleObjectsReturned
diff --git a/tortoise/exceptions.py b/tortoise/exceptions.py
index bc0ca37e4..1ccdb51d7 100644
--- a/tortoise/exceptions.py
+++ b/tortoise/exceptions.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Union, Optional
 
 if TYPE_CHECKING:
     from tortoise import Model, Type
@@ -52,18 +52,30 @@ class NoValuesFetched(OperationalError):
     """
 
 
-class MultipleObjectsReturned(OperationalError):
+class NotExistOrMultiple(OperationalError):
+    TEMPLATE = ""
+
+    def __init__(self, model: "Union[Type[Model], str]", *args) -> None:
+        self.model: "Optional[Type[Model]]" = None
+        if isinstance(model, str):
+            args = (model,) + args
+        else:
+            self.model = model
+        super().__init__(*args)
+
+    def __str__(self) -> str:
+        if self.model is None:
+            return super().__str__()
+        return self.TEMPLATE.format(self.model.__name__)
+
+
+class MultipleObjectsReturned(NotExistOrMultiple):
     """
     The MultipleObjectsReturned exception is raised when doing a ``.get()`` operation,
     and more than one object is returned.
     """
 
-    def __init__(self, model: "Type[Model]", *args):
-        self.model: "Type[Model]" = model
-        super().__init__(*args)
-
-    def __str__(self):
-        return f'Multiple objects returned for "{self.model.__name__}", expected exactly one'
+    TEMPLATE = 'Multiple objects returned for "{}", expected exactly one'
 
 
 class ObjectDoesNotExistError(OperationalError, KeyError):
@@ -80,17 +92,12 @@ def __str__(self):
         return f"{self.model.__name__} has no object with {self.pk_name}={self.pk_val}"
 
 
-class DoesNotExist(OperationalError):
+class DoesNotExist(NotExistOrMultiple):
     """
     The DoesNotExist exception is raised when expecting data, such as a ``.get()`` operation.
     """
 
-    def __init__(self, model: "Type[Model]", *args):
-        self.model: "Type[Model]" = model
-        super().__init__(*args)
-
-    def __str__(self):
-        return f'Object "{self.model.__name__}" does not exist'
+    TEMPLATE = 'Object "{}" does not exist'
 
 
 class IncompleteInstanceError(OperationalError):