<Test>: Fix test case

app/routers/video.py

@@ -24,7 +24,7 @@ from app.routers.image import inferenceImage
 from google.cloud.firestore_v1.base_query import FieldFilter
 from google.cloud.firestore import ArrayUnion
 from app import logger
-
+from custom_utils import utils
 router = APIRouter(prefix="/video", tags=["Video"])
 
 
@@ -166,51 +166,4 @@ def updateArtifact(artifactId: str, body):
     artifact_snapshot = artifact_ref.get()
     if artifact_snapshot.exists:
         artifact_ref.update(body)
-    # sendMessage(artifactId)
-
-
-# This function cannot be automation test because the requirement of another device to receive notification
-def sendMessage(artifactId: str, message: str = None):
-    token = []
-    artifact = db.collection("artifacts").document(artifactId).get()
-    if not artifact.exists:
-        return
-    user_ref = db.collection("user").where(
-        filter=FieldFilter("artifacts", "array-contains", "artifacts/" + artifactId)
-    )
-    for user in user_ref:
-        token.append(user.get().to_dict()["deviceId"])
-    if message is not None:
-        messaging.MulticastMessage(
-            data={
-                "notification": {
-                    "title": message,
-                    "body": "Video "
-                    + artifact.name
-                    + " has done inference. Click here to see the video",
-                },
-            },
-            android=messaging.AndroidConfig(
-                notification=messaging.AndroidNotification(
-                    icon="stock_ticker_update", color="#f45342"
-                ),
-            ),
-        )
-    else:
-        messaging.MulticastMessage(
-            data={
-                "notification": {
-                    "title": "Video " + artifact.name + " has done inference.",
-                    "body": "Video "
-                    + artifact.name
-                    + " has done inference. Click here to see the video",
-                },
-            },
-            android=messaging.AndroidConfig(
-                notification=messaging.AndroidNotification(
-                    icon="stock_ticker_update", color="#f45342"
-                ),
-            ),
-        )
-    response = messaging.send_multicast(message)
-    return response.success_count
+    utils.sendMessage(artifactId)

custom_utils/utils.py

@@ -0,0 +1,52 @@
+from firebase_admin import messaging
+from app import db
+from google.cloud.firestore_v1.base_query import FieldFilter
+
+# This function cannot be automation test because the requirement of another device to receive notification
+def sendMessage(artifactId: str, message: str = None):
+    token = []
+    artifact = db.collection("artifacts").document(artifactId).get()
+    if not artifact.exists:
+        return
+    artifact = artifact.to_dict()
+    user_ref = db.collection("user").where(
+        filter=FieldFilter("artifacts", "array_contains", "artifacts/" + artifactId)
+    ).stream()
+    for user in user_ref:
+        token.append(user.to_dict()["deviceId"])
+    if message is not None:
+        msg = messaging.MulticastMessage(
+            data={
+                "notification": {
+                    "title": message,
+                    "body": "Video "
+                    + artifact["name"]
+                    + " has done inference. Click here to see the video",
+                },
+            },
+            android=messaging.AndroidConfig(
+                notification=messaging.AndroidNotification(
+                    icon="stock_ticker_update", color="#f45342"
+                ),
+            ),
+            tokens=token
+        )
+    else:
+        msg = messaging.MulticastMessage(
+            data={
+                "notification": {
+                    "title": "Video " + artifact['name'] + " has done inference.",
+                    "body": "Video "
+                    + artifact["name"]
+                    + " has done inference. Click here to see the video",
+                },
+            },
+            android=messaging.AndroidConfig(
+                notification=messaging.AndroidNotification(
+                    icon="stock_ticker_update", color="#f45342"
+                ),
+            ),
+            tokens=token
+        )
+    response = messaging.send_multicast(msg)
+    return response.success_count

tests/test_friend_request.py

@@ -11,7 +11,7 @@ from app.constants import deviceId
 from fastapi.routing import APIRoute
 from app import db
 from google.cloud.firestore_v1.base_query import FieldFilter
-
+import time
 
 def endpoints():
     endpoints = []
@@ -235,10 +235,17 @@ class TestFriendRequest:
             "Content-Type": "application/json",
             "Authorization": "Bearer " + inviter["token"],
         }
-        response = client.request(
-            "POST", "friend_request", headers=headers, data=payload
-        )
-        assert response.status_code == 200
+        flag = False
+        for i in range(5):
+            response = client.request(
+                "POST", "friend_request", headers=headers, data=payload
+            )
+            if(response.status_code == 200):
+                flag = True
+                break
+            else:
+                time.sleep(2)
+        assert flag == True
         result = mmcv.imfrombytes(response.read())
         # Check returned QR image
         assert result.shape[2] == 3

tests/test_video.py

@@ -101,7 +101,8 @@ class TestVideoAPI:
         for doc in docs:
             # For each document in docs. Verify name and status of the artifact
             index += 1
-            data = doc.get().to_dict()
+            data = doc.to_dict()
+            print(data)
             assert data["name"] == artifactName
             assert data["status"] == "pending"
             assert index == 1