diff --git a/tests/unit/steps/argilla/test_preference.py b/tests/unit/steps/argilla/test_preference.py
index 5ff95782b0..edeb566849 100644
--- a/tests/unit/steps/argilla/test_preference.py
+++ b/tests/unit/steps/argilla/test_preference.py
@@ -15,7 +15,7 @@
 import os
 from unittest.mock import patch
 
-import argilla_sdk as rg
+import argilla as rg
 from distilabel.pipeline.local import Pipeline
 from distilabel.steps.argilla.preference import PreferenceToArgilla
 
diff --git a/tests/unit/steps/argilla/test_text_generation.py b/tests/unit/steps/argilla/test_text_generation.py
index 5679c54517..884ae51cff 100644
--- a/tests/unit/steps/argilla/test_text_generation.py
+++ b/tests/unit/steps/argilla/test_text_generation.py
@@ -15,7 +15,7 @@
 import os
 from unittest.mock import patch
 
-import argilla_sdk as rg
+import argilla as rg
 from distilabel.pipeline.local import Pipeline
 from distilabel.steps.argilla.text_generation import TextGenerationToArgilla