diff --git a/app/cafe.py b/app/cafe.py
index 351e1d60..d52725ff 100644
--- a/app/cafe.py
+++ b/app/cafe.py
@@ -1,4 +1,5 @@
 import datetime
+
 from app.errors import (
     NotVaccinatedError,
     OutdatedVaccineError,
diff --git a/app/main.py b/app/main.py
index 3d171dbd..0be02e40 100644
--- a/app/main.py
+++ b/app/main.py
@@ -3,20 +3,25 @@
 from app.cafe import Cafe
 
 
-def go_to_cafe(friends: list, cafe: Cafe) -> str | VaccineError:
-    masks_to_buy = 0
+def get_amount_without_mask(friends: list) -> int:
+    return len([friend for friend in friends if not friend["wearing_a_mask"]])
 
-    for friend in friends:
+
+def go_to_cafe(friends: list, cafe: Cafe) -> str:
+
+    for i, friend in enumerate(friends):
         try:
             cafe.visit_cafe(friend)
+            if i == len(friends) - 1 and get_amount_without_mask(friends) != 0:
+                return (f"Friends should buy "
+                        f"{get_amount_without_mask(friends)} masks")
 
         except VaccineError:
             return "All friends should be vaccinated"
 
         except NotWearingMaskError:
-            masks_to_buy += 1
-
-    if masks_to_buy != 0:
-        return f"Friends should buy {masks_to_buy} masks"
+            if i == len(friends) - 1 and get_amount_without_mask(friends) != 0:
+                return (f"Friends should buy "
+                        f"{get_amount_without_mask(friends)} masks")
 
     return f"Friends can go to {cafe.name}"