From 01eecf901f5a3d2521b4c2d349fd95629fc1e5f5 Mon Sep 17 00:00:00 2001
From: Abhyudai <13880786+abhiabhi94@users.noreply.github.com>
Date: Sun, 1 Aug 2021 13:42:06 +0530
Subject: [PATCH] Add tests for some queryset methods (#684)

---
 .../managers/querysets/test_basic_methods.yml | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/tests/typecheck/managers/querysets/test_basic_methods.yml b/tests/typecheck/managers/querysets/test_basic_methods.yml
index aec7a53..8a8d078 100644
--- a/tests/typecheck/managers/querysets/test_basic_methods.yml
+++ b/tests/typecheck/managers/querysets/test_basic_methods.yml
@@ -1,5 +1,6 @@
 -   case: queryset_basic_methods_return_type
     main: |
+        from django.utils import timezone
         from myapp.models import Blog
 
         qs = Blog.objects.all()
@@ -11,14 +12,27 @@
         reveal_type(qs.earliest()) # N: Revealed type is "myapp.models.Blog*"
         reveal_type(qs[0]) # N: Revealed type is "myapp.models.Blog*"
         reveal_type(qs[:9]) # N: Revealed type is "django.db.models.query._QuerySet[myapp.models.Blog, myapp.models.Blog]"
-        reveal_type(qs.in_bulk()) # N: Revealed type is "builtins.dict[Any, myapp.models.Blog*]"
-
+        reveal_type(qs.create()) # N: Revealed type is "myapp.models.Blog*"
+        reveal_type(qs.get_or_create()) # N: Revealed type is "Tuple[myapp.models.Blog*, builtins.bool]"
+        reveal_type(qs.exists()) # N: Revealed type is "builtins.bool"
+        reveal_type(qs.none()) # N: Revealed type is "django.db.models.query._QuerySet[myapp.models.Blog, myapp.models.Blog]"
+        reveal_type(qs.update_or_create()) # N: Revealed type is "Tuple[myapp.models.Blog*, builtins.bool]"
+        reveal_type(qs.explain()) # N: Revealed type is "builtins.str"
+        reveal_type(qs.raw(qs.explain())) # N: Revealed type is "django.db.models.query.RawQuerySet[Any]"
         # .dates / .datetimes
         reveal_type(Blog.objects.dates("created_at", "day")) # N: Revealed type is "django.db.models.query._QuerySet[myapp.models.Blog*, datetime.date]"
         reveal_type(Blog.objects.datetimes("created_at", "day")) # N: Revealed type is "django.db.models.query._QuerySet[myapp.models.Blog*, datetime.datetime]"
 
         # AND-ing QuerySets
         reveal_type(Blog.objects.all() & Blog.objects.all()) # N: Revealed type is "django.db.models.query._QuerySet[myapp.models.Blog, myapp.models.Blog]"
+
+        # bulk methods
+        reveal_type(qs.count()) # N: Revealed type is "builtins.int"
+        reveal_type(qs.update(created_at=timezone.now()))  # N: Revealed type is "builtins.int"
+        reveal_type(qs.in_bulk()) # N: Revealed type is "builtins.dict[Any, myapp.models.Blog*]"
+        reveal_type(qs.bulk_update(list(qs), fields=["created_at"]))  # N: Revealed type is "None"
+        reveal_type(qs.bulk_create([]))  # N: Revealed type is "builtins.list[myapp.models.Blog*]"
+        reveal_type(qs.delete()) # N: Revealed type is "Tuple[builtins.int, builtins.dict[builtins.str, builtins.int]]"
     installed_apps:
         - myapp
     files:
@@ -30,7 +44,6 @@
                 class Blog(models.Model):
                     created_at = models.DateTimeField()
 
-
 -   case: queryset_missing_method
     main: |
         from myapp.models import User