typing

demml · May 6, 2024 · c819046 · c819046
1 parent 79fcfcb
commit c819046
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 17 deletions.
diff --git a/opsml/app/routes/cards.py b/opsml/app/routes/cards.py
@@ -101,7 +101,7 @@ def query_registry_stats(
 
     try:
         registry: CardRegistry = getattr(request.app.state.registries, registry_type)
-        stats = registry._registry.query_stats(search_term)
+        stats: Dict[str, int] = registry._registry.query_stats(search_term)
 
         return stats
 
@@ -154,8 +154,14 @@ def query_registry_page(
             FastAPI request object
         registry_type:
             Type of registry
-        uid:
-            uid of the card
+        sort_by:
+            Field to sort by
+        repository:
+            repository to filter by
+        search_term:
+            search term to filter by. This term can be a repository or a name
+        page:
+            page number
 
     Returns:
         `dict`

diff --git a/opsml/registry/sql/base/query_engine.py b/opsml/registry/sql/base/query_engine.py
@@ -1,4 +1,5 @@
 # mypy: disable-error-code="call-overload"
+# pylint: disable=not-callable
 
 # Copyright (c) Shipt, Inc.
 # This source code is licensed under the MIT license found in the
@@ -444,10 +445,8 @@ def query_stats(
         Args:
             table:
                 Registry table to query
-            repository:
-                Repository name
-            name:
-                Name of the card
+            search_term:
+                Search term
         """
 
         query = select(
@@ -491,10 +490,12 @@ def query_page(
         Args:
             sort_by:
                 Field to sort by
-            name:
-                Name of the card
+            page:
+                Page number
+            search_term:
+                Search term
             repository:
-                Repository of the card
+                Repository name
             table:
                 Registry table to query
         Returns:

diff --git a/opsml/registry/sql/base/server.py b/opsml/registry/sql/base/server.py
@@ -75,10 +75,9 @@ def unique_repositories(self) -> Sequence[str]:
     def query_stats(self, search_term: Optional[str] = None) -> Dict[str, int]:
         """Query stats from Card Database
         Args:
-            repository:
-                Repository to filter by
-            name:
-                Card name to filter by
+            search_term:
+                Search term to filter by
+
         Returns:
             Dictionary of stats
         """
@@ -95,10 +94,12 @@ def query_page(
         Args:
             sort_by:
                 Field to sort by
+            page:
+                Page number
             repository:
                 Repository to filter by
-            name:
-                Card name to filter by
+            search_term:
+                Search term to filter by
         Returns:
             List of tuples
         """

diff --git a/tests/test_ui/test_ui_routes.py b/tests/test_ui/test_ui_routes.py
@@ -33,9 +33,20 @@ def test_card_routes(
 
     res = response.json()
     assert response.status_code == 200
-
     assert res == {"nbr_names": 1, "nbr_versions": 1, "nbr_repos": 1}
 
+    # force error
+    response = test_app.get(
+        url="/opsml/cards/registry/stats",
+        params={
+            "registry_type": "error",
+            "search_term": modelcard.repository,
+        },
+    )
+
+    res = response.json()
+    assert response.status_code == 500
+
     response = test_app.get(
         url="/opsml/cards/registry/query/page",
         params={
@@ -49,3 +60,14 @@ def test_card_routes(
     assert response.status_code == 200
 
     assert len(res["page"]) == 1
+
+    # force error
+    response = test_app.get(
+        url="/opsml/cards/registry/query/page",
+        params={
+            "registry_type": "error",
+            "search_term": modelcard.name,
+            "repository": modelcard.repository,
+        },
+    )
+    assert response.status_code == 500