Checkpoint

bookwyrm/book_search.py

@@ -43,6 +43,7 @@ def search(
     min_confidence: float = 0,
     filters: Optional[list[Any]] = None,
     return_first: bool = False,
+    books = None
 ) -> Union[Optional[models.Edition], QuerySet[models.Edition]]:
     """search your local database"""
     filters = filters or []
@@ -54,17 +55,16 @@ def search(
     # first, try searching unique identifiers
     # unique identifiers never have spaces, title/author usually do
     if not " " in query:
-        results = search_identifiers(query, *filters, return_first=return_first)
+        results = search_identifiers(query, *filters, return_first=return_first, books=books)
 
     # if there were no identifier results...
     if not results:
         # then try searching title/author
         results = search_title_author(
-            query, min_confidence, *filters, return_first=return_first
+            query, min_confidence, *filters, return_first=return_first, books=books
         )
     return results
 
-
 def isbn_search(query):
     """search your local database"""
     if not query:
@@ -98,8 +98,9 @@ def format_search_result(search_result):
 
 
 def search_identifiers(
-    query, *filters, return_first=False
+    query, *filters, return_first=False, books=None,
 ) -> Union[Optional[models.Edition], QuerySet[models.Edition]]:
+    books = books or models.Edition
     """tries remote_id, isbn; defined as dedupe fields on the model"""
     if connectors.maybe_isbn(query):
         # Oh did you think the 'S' in ISBN stood for 'standard'?
@@ -111,7 +112,7 @@ def search_identifiers(
         for f in models.Edition._meta.get_fields()
         if hasattr(f, "deduplication_field") and f.deduplication_field
     ]
-    results = models.Edition.objects.filter(
+    results = books.filter(
         *filters, reduce(operator.or_, (Q(**f) for f in or_filters))
     ).distinct()
 
@@ -121,12 +122,17 @@ def search_identifiers(
 
 
 def search_title_author(
-    query, min_confidence, *filters, return_first=False
+    query,
+    min_confidence,
+    *filters,
+    return_first=False,
+    books=None,
 ) -> QuerySet[models.Edition]:
     """searches for title and author"""
+    books = books or models.Edition.objects
     query = SearchQuery(query, config="simple") | SearchQuery(query, config="english")
     results = (
-        models.Edition.objects.filter(*filters, search_vector=query)
+        books.filter(*filters, search_vector=query)
         .annotate(rank=SearchRank(F("search_vector"), query))
         .filter(rank__gt=min_confidence)
         .order_by("-rank")