diff --git a/tests/unit_tests.py b/tests/unit_tests.py
index a830aa0..e6a1b2b 100644
--- a/tests/unit_tests.py
+++ b/tests/unit_tests.py
@@ -1,4 +1,5 @@
 from yml2block.__main__ import guess_input_type
+from yml2block.rules import Level
 
 
 def test_input_guessing_valid_tsv():
@@ -48,7 +49,7 @@ def test_input_guessing_valid_csv():
     for path in valid_csv_paths:
         guessed_type, violations = guess_input_type(path)
         assert len(violations) == 1
-        assert violations[0].level == "WARNING"
+        assert violations[0].level == Level.WARNING
         assert guessed_type == "csv"
 
 
@@ -114,5 +115,5 @@ def test_input_guessing_invalid_extension():
     for path in invalid_extension_paths:
         guessed_type, violations = guess_input_type(path)
         assert len(violations) == 1
-        assert violations[0].level == "ERROR"
+        assert violations[0].level == Level.ERROR
         assert guessed_type == False