diff --git a/tests/test_graph.py b/tests/test_graph.py
index 38ec118..4ca84ce 100644
--- a/tests/test_graph.py
+++ b/tests/test_graph.py
@@ -1,15 +1,19 @@
 import pytest
-from package.graph import display,copy_to_clipboard
+from package.graph import display
 from datetime import datetime
 import subprocess
+from package.graph import copy_to_clipboard
+import os
+import sys
+import shutil
 
 def test_display_right_password():
     current_time = datetime.now().strftime("%H%M")
-    assert display("test",current_time) is None 
+    assert display("test",current_time) == None 
 
 def test_display_wrong_password():
-    with pytest.raises(ValueError, match="syntax error"):
-        display("test",1234)
+    with pytest.raises(ValueError, match="Invalid password"):
+        display("test", 1111)
 
 def test_copy_to_clipboard(monkeypatch):
     def mock_run(*args, **kwargs):
@@ -19,4 +23,4 @@ def mock_run(*args, **kwargs):
     try: 
         copy_to_clipboard("Whatever test")
     except Exception as e:
-        pytest.fail(f"Unexpected error raised: {e}")
\ No newline at end of file
+        pytest.fail("Unexpected error raised: {e}")
\ No newline at end of file