diff --git a/tests/test_main.py b/tests/test_main.py
index 19d7abf8..ab49b418 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -1,7 +1,9 @@
-import pytest
 import inspect
 
-from app import main
+import pytest
+
+from app.main import count_occurrences
+
 
 @pytest.mark.parametrize(
     "phrase,letter,count",
@@ -12,16 +14,17 @@
         ("Abracadabra", "A", 5),
         ("", "a", 0),
         ("Samsung", "b", 0),
-    ]
+    ],
 )
 def test_count_occurrences(phrase, letter, count):
-    assert main.count_occurrences(phrase, letter) == count, (
+    assert count_occurrences(phrase, letter) == count, (
         f"Function 'count_occurrences' should return {count}, "
         f"when 'phrase'='{phrase}' and 'letter'='{letter}'"
     )
 
 
 def test_removed_comment():
-    lines = inspect.getsource(main)
-    assert "# write your code here" not in lines, ("You have to"
-        " remove the unnecessary comment '# write your code here'")
+    lines = inspect.getsource(count_occurrences)
+    assert "# write your code here" not in lines, (
+        "You have to" " remove the unnecessary comment '# write your code here'"
+    )