From 885b24b35b7662f3073839ccb9d5b470696b6d85 Mon Sep 17 00:00:00 2001
From: Simon Prickett <simon@redislabs.com>
Date: Mon, 14 Feb 2022 20:44:49 +0000
Subject: [PATCH] Adds delete function at a class level.

---
 aredis_om/model/model.py |  8 ++++++++
 tests/test_hash_model.py | 15 +++++++++++++++
 tests/test_json_model.py | 15 +++++++++++++++
 3 files changed, 38 insertions(+)

diff --git a/aredis_om/model/model.py b/aredis_om/model/model.py
index 8a1fac2..0d9ea05 100644
--- a/aredis_om/model/model.py
+++ b/aredis_om/model/model.py
@@ -1335,6 +1335,10 @@ class HashModel(RedisModel, abc.ABC):
             result = cls.parse_obj(document)
         return result
 
+    @classmethod
+    async def delete(cls, pk: Any) -> int:
+        return await cls.db().delete(cls.make_primary_key(pk))
+
     @classmethod
     @no_type_check
     def _get_value(cls, *args, **kwargs) -> Any:
@@ -1503,6 +1507,10 @@ class JsonModel(RedisModel, abc.ABC):
             raise NotFoundError
         return cls.parse_raw(document)
 
+    @classmethod
+    async def delete(cls, pk: Any) -> int:
+        return await cls.db().delete(cls.make_primary_key(pk))
+
     @classmethod
     def redisearch_schema(cls):
         key_prefix = cls.make_key(cls._meta.primary_key_pattern.format(pk=""))
diff --git a/tests/test_hash_model.py b/tests/test_hash_model.py
index 9a9291a..a48999f 100644
--- a/tests/test_hash_model.py
+++ b/tests/test_hash_model.py
@@ -374,6 +374,21 @@ async def test_saves_model_and_creates_pk(m):
     member2 = await m.Member.get(member.pk)
     assert member2 == member
 
+@pytest.mark.asyncio
+async def test_delete(m):
+    member = m.Member(
+        first_name="Simon",
+        last_name="Prickett",
+        email="s@example.com",
+        join_date=today,
+        age=97,
+        bio="This is a test use to be deleted.",
+    )
+
+    await member.save()
+    response = await m.Member.delete(member.pk)
+    assert response == 1
+
 
 def test_raises_error_with_embedded_models(m):
     class Address(m.BaseHashModel):
diff --git a/tests/test_json_model.py b/tests/test_json_model.py
index 8b360f7..ce240e0 100644
--- a/tests/test_json_model.py
+++ b/tests/test_json_model.py
@@ -186,6 +186,21 @@ async def test_saves_model_and_creates_pk(address, m, redis):
     assert member2 == member
     assert member2.address == address
 
+@pytest.mark.asyncio
+async def test_delete(address, m, redis):
+    member = m.Member(
+        first_name="Simon",
+        last_name="Prickett",
+        email="s@example.com",
+        join_date=today,
+        age=38,
+        address=address,
+    )
+
+    await member.save()
+    response = await m.Member.delete(member.pk)
+    assert response == 1
+
 
 @pytest.mark.asyncio
 async def test_saves_many_implicit_pipeline(address, m):