diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 6130860..48a0644 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -1,6 +1,7 @@
 import pytest
 
 from pyredis.protocol import extract_frame_from_buffer, encode_message
+from pyredis.datastore import DataStore
 from pyredis.types import (
     Array,
     BulkString,