add endpoint for getting monthly workflow (#141)

* add endpoint for getting monthly workflow

Author: tianj7

src/argowrapper/routes/routes.py

@@ -252,11 +252,10 @@ def check_user_info_for_billing_id_and_workflow_limit(request):
         return None, None
 
 
-def check_user_reached_monthly_workflow_cap(request_token, billing_id, custom_limit):
+def check_user_monthly_workflow_cap(request_token, billing_id, custom_limit):
     """
-    Query Argo service to see how many successful run user already
-    have in the current calendar month. If the number is greater than
-    the threshold, return error.
+    Query Argo service to see how many workflow runs user already
+    have in the current calendar month. Return number of workflow runs and limit
     """
 
     try:
@@ -271,21 +270,7 @@ def check_user_reached_monthly_workflow_cap(request_token, billing_id, custom_li
                 limit = GEN3_NON_VA_WORKFLOW_MONTHLY_CAP
             else:
                 limit = GEN3_DEFAULT_WORKFLOW_MONTHLY_CAP
-
-        if len(current_month_workflows) >= limit:
-            logger.warn(
-                "This user {} already executed {} workflows this month and cannot create new ones anymore. The currently monthly cap for this user is {}.".format(
-                    username, len(current_month_workflows), limit
-                )
-            )
-            return True
-        else:
-            logger.info(
-                "This user {} executed {} workflows this month. The currently monthly cap for this user is {}.".format(
-                    username, len(current_month_workflows), limit
-                )
-            )
-        return False
+        return len(current_month_workflows), limit
     except Exception as e:
         logger.error(e)
         traceback.print_exc()
@@ -316,10 +301,12 @@ def submit_workflow(
         )
 
         # if user has billing_id (non-VA user), check if they already reached the monthly cap
-        reached_monthly_cap = check_user_reached_monthly_workflow_cap(
+        workflow_run, workflow_limit = check_user_monthly_workflow_cap(
             request.headers.get("Authorization"), billing_id, workflow_limit
         )
 
+        reached_monthly_cap = workflow_run >= workflow_limit
+
         # submit workflow:
         if not reached_monthly_cap:
             return argo_engine.workflow_submission(
@@ -448,3 +435,30 @@ def get_workflow_logs(
             content="Unexpected Error Occurred",
             status_code=HTTP_500_INTERNAL_SERVER_ERROR,
         )
+
+
+@router.get("/workflows/user-monthly", status_code=HTTP_200_OK)
+def get_user_monthly_workflow(
+    request: Request,
+) -> Dict[str, any]:
+    """
+    Query Argo service to see how many successful run user already
+    have in the current calendar month. Return workflow numbers and workflow cap
+    """
+
+    try:
+        billing_id, workflow_limit = check_user_info_for_billing_id_and_workflow_limit(
+            request
+        )
+
+        # if user has billing_id (non-VA user), check if they already reached the monthly cap
+        workflow_run, workflow_limit = check_user_monthly_workflow_cap(
+            request.headers.get("Authorization"), billing_id, workflow_limit
+        )
+
+        result = {"workflow_run": workflow_run, "workflow_limit": workflow_limit}
+        return result
+    except Exception as e:
+        logger.error(e)
+        traceback.print_exc()
+        raise e

test/test_routes.py

@@ -8,8 +8,15 @@
 from fastapi.testclient import TestClient
 from argowrapper.constants import *
 from test.constants import EXAMPLE_AUTH_HEADER
-from argowrapper.routes.routes import router, check_user_reached_monthly_workflow_cap
-from argowrapper.constants import GEN3_NON_VA_WORKFLOW_MONTHLY_CAP
+from argowrapper.routes.routes import (
+    router,
+    check_user_monthly_workflow_cap,
+    get_user_monthly_workflow,
+)
+from argowrapper.constants import (
+    GEN3_NON_VA_WORKFLOW_MONTHLY_CAP,
+    GEN3_DEFAULT_WORKFLOW_MONTHLY_CAP,
+)
 
 variables = [
     {"variable_type": "concept", "concept_id": "2000000324"},
@@ -110,13 +117,13 @@ def test_submit_workflow(client):
     ) as mock_check_billing_id, patch(
         "requests.get"
     ) as mock_requests, patch(
-        "argowrapper.routes.routes.check_user_reached_monthly_workflow_cap"
+        "argowrapper.routes.routes.check_user_monthly_workflow_cap"
     ) as mock_check_monthly_cap:
         mock_auth.return_value = True
         mock_engine.return_value = "workflow_123"
         mock_check_billing_id.return_value = None, None
         mock_requests.side_effect = mocked_requests_get
-        mock_check_monthly_cap.return_value = False
+        mock_check_monthly_cap.return_value = 1, 3
 
         response = client.post(
             "/submit",
@@ -558,9 +565,9 @@ def test_submit_workflow_with_user_billing_id(client):
         data["user_tags"] = {"tags": {"othertag1": "tag1", "billing_id": "1234"}}
 
         with patch(
-            "argowrapper.routes.routes.check_user_reached_monthly_workflow_cap"
+            "argowrapper.routes.routes.check_user_monthly_workflow_cap"
         ) as mock_check_monthly_cap:
-            mock_check_monthly_cap.return_value = False
+            mock_check_monthly_cap.return_value = 1, 3
             response = client.post(
                 "/submit",
                 data=json.dumps(data),
@@ -584,7 +591,7 @@ def test_submit_workflow_with_user_billing_id(client):
         assert response.status_code == 500
 
 
-def test_check_user_reached_monthly_workflow_cap():
+def test_check_user_monthly_workflow_cap():
     headers = {
         "Content-Type": "application/json",
         "Authorization": EXAMPLE_AUTH_HEADER,
@@ -600,16 +607,14 @@ def test_check_user_reached_monthly_workflow_cap():
 
         # Test Under Default Limit
         assert (
-            check_user_reached_monthly_workflow_cap(
-                headers["Authorization"], None, None
-            )
-            == False
+            check_user_monthly_workflow_cap(headers["Authorization"], None, None) == 2,
+            GEN3_NON_VA_WORKFLOW_MONTHLY_CAP,
         )
 
         # Test Custom Limit
         assert (
-            check_user_reached_monthly_workflow_cap(headers["Authorization"], None, 2)
-            == True
+            check_user_monthly_workflow_cap(headers["Authorization"], None, 2) == 2,
+            2,
         )
 
         # Test Billing Id User Exceeding Limit
@@ -619,10 +624,9 @@ def test_check_user_reached_monthly_workflow_cap():
         mock_get_workflow.return_value = workflows
 
         assert (
-            check_user_reached_monthly_workflow_cap(
-                headers["Authorization"], "1234", None
-            )
-            == True
+            check_user_monthly_workflow_cap(headers["Authorization"], "1234", None)
+            == GEN3_NON_VA_WORKFLOW_MONTHLY_CAP + 1,
+            GEN3_NON_VA_WORKFLOW_MONTHLY_CAP,
         )
 
         # Test VA User Exceeding Limit
@@ -631,10 +635,9 @@ def test_check_user_reached_monthly_workflow_cap():
             workflows.append({"wf_name": "workflow" + str(index)})
         mock_get_workflow.return_value = workflows
         assert (
-            check_user_reached_monthly_workflow_cap(
-                headers["Authorization"], None, None
-            )
-            == True
+            check_user_monthly_workflow_cap(headers["Authorization"], None, None)
+            == GEN3_DEFAULT_WORKFLOW_MONTHLY_CAP + 1,
+            GEN3_DEFAULT_WORKFLOW_MONTHLY_CAP,
         )
 
 
@@ -646,14 +649,14 @@ def test_submit_workflow_with_billing_id_and_over_monthly_cap(client):
     ) as mock_log, patch(
         "argowrapper.routes.routes.check_user_info_for_billing_id_and_workflow_limit"
     ) as mock_check_billing_id, patch(
-        "argowrapper.routes.routes.check_user_reached_monthly_workflow_cap"
+        "argowrapper.routes.routes.check_user_monthly_workflow_cap"
     ) as mock_check_monthly_cap, patch(
         "requests.get"
     ) as mock_requests:
         mock_auth.return_value = True
         mock_engine.return_value = "workflow_123"
         mock_check_billing_id.return_value = "1234", None
-        mock_check_monthly_cap.return_value = True
+        mock_check_monthly_cap.return_value = 1, 1
         mock_requests.side_effect = mocked_requests_get
 
         response = client.post(
@@ -675,14 +678,14 @@ def test_submit_workflow_over_monthly_cap(client):
     ) as mock_log, patch(
         "argowrapper.routes.routes.check_user_info_for_billing_id_and_workflow_limit"
     ) as mock_check_billing_id, patch(
-        "argowrapper.routes.routes.check_user_reached_monthly_workflow_cap"
+        "argowrapper.routes.routes.check_user_monthly_workflow_cap"
     ) as mock_check_monthly_cap, patch(
         "requests.get"
     ) as mock_requests:
         mock_auth.return_value = True
         mock_engine.return_value = "workflow_123"
         mock_check_billing_id.return_value = None, None
-        mock_check_monthly_cap.return_value = True
+        mock_check_monthly_cap.return_value = 1, 1
         mock_requests.side_effect = mocked_requests_get
 
         response = client.post(
@@ -722,3 +725,52 @@ def test_submit_workflow_with_non_team_project_cohort(client):
             },
         )
         assert response.status_code == 400
+
+
+def test_get_user_monthly_workflow(client):
+    with patch(
+        "argowrapper.engine.argo_engine.ArgoEngine.get_user_workflows_for_current_month"
+    ) as mock_get_workflow, patch(
+        "argowrapper.routes.routes.check_user_info_for_billing_id_and_workflow_limit"
+    ) as mock_check_billing_id:
+        mock_get_workflow.return_value = [
+            {"wf_name": "workflow1"},
+            {"wf_name": "workflow2"},
+        ]
+        mock_check_billing_id.return_value = None, None
+
+        response = client.get(
+            "/workflows/user-monthly",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": EXAMPLE_AUTH_HEADER,
+            },
+        )
+
+        assert response.status_code == 200
+        assert response.content.decode(
+            "utf-8"
+        ) == '{{"workflow_run":2,"workflow_limit":{}}}'.format(
+            GEN3_DEFAULT_WORKFLOW_MONTHLY_CAP
+        )
+
+        mock_get_workflow.return_value = [
+            {"wf_name": "workflow1"},
+            {"wf_name": "workflow2"},
+            {"wf_name": "workflow3"},
+        ]
+        mock_check_billing_id.return_value = "1234", None
+
+        response = client.get(
+            "/workflows/user-monthly",
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": EXAMPLE_AUTH_HEADER,
+            },
+        )
+
+        assert response.content.decode(
+            "utf-8"
+        ) == '{{"workflow_run":3,"workflow_limit":{}}}'.format(
+            GEN3_NON_VA_WORKFLOW_MONTHLY_CAP
+        )