diff --git a/kedro-airflow/features/steps/cli_steps.py b/kedro-airflow/features/steps/cli_steps.py
index 23eb58727..7bfa482ac 100644
--- a/kedro-airflow/features/steps/cli_steps.py
+++ b/kedro-airflow/features/steps/cli_steps.py
@@ -20,27 +20,27 @@ def init_airflow(context, home_dir):
 def prepare_old_catalog(context):
     config = {
         "example_train_x": {
-            "type": "PickleLocalDataSet",
+            "type": "PickleLocalDataset",
             "filepath": "data/02_intermediate/example_train_x.pkl",
         },
         "example_train_y": {
-            "type": "PickleLocalDataSet",
+            "type": "PickleLocalDataset",
             "filepath": "data/02_intermediate/example_train_y.pkl",
         },
         "example_test_x": {
-            "type": "PickleLocalDataSet",
+            "type": "PickleLocalDataset",
             "filepath": "data/02_intermediate/example_test_x.pkl",
         },
         "example_test_y": {
-            "type": "PickleLocalDataSet",
+            "type": "PickleLocalDataset",
             "filepath": "data/02_intermediate/example_test_y.pkl",
         },
         "example_model": {
-            "type": "PickleLocalDataSet",
+            "type": "PickleLocalDataset",
             "filepath": "data/02_intermediate/example_model.pkl",
         },
         "example_predictions": {
-            "type": "PickleLocalDataSet",
+            "type": "PickleLocalDataset",
             "filepath": "data/02_intermediate/example_predictions.pkl",
         },
     }
@@ -53,27 +53,27 @@ def prepare_old_catalog(context):
 def prepare_catalog(context):
     config = {
         "example_train_x": {
-            "type": "pickle.PickleDataSet",
+            "type": "pickle.PickleDataset",
             "filepath": "data/02_intermediate/example_train_x.pkl",
         },
         "example_train_y": {
-            "type": "pickle.PickleDataSet",
+            "type": "pickle.PickleDataset",
             "filepath": "data/02_intermediate/example_train_y.pkl",
         },
         "example_test_x": {
-            "type": "pickle.PickleDataSet",
+            "type": "pickle.PickleDataset",
             "filepath": "data/02_intermediate/example_test_x.pkl",
         },
         "example_test_y": {
-            "type": "pickle.PickleDataSet",
+            "type": "pickle.PickleDataset",
             "filepath": "data/02_intermediate/example_test_y.pkl",
         },
         "example_model": {
-            "type": "pickle.PickleDataSet",
+            "type": "pickle.PickleDataset",
             "filepath": "data/02_intermediate/example_model.pkl",
         },
         "example_predictions": {
-            "type": "pickle.PickleDataSet",
+            "type": "pickle.PickleDataset",
             "filepath": "data/02_intermediate/example_predictions.pkl",
         },
     }
diff --git a/kedro-telemetry/tests/test_plugin.py b/kedro-telemetry/tests/test_plugin.py
index ccbaf8afe..bee020e7d 100644
--- a/kedro-telemetry/tests/test_plugin.py
+++ b/kedro-telemetry/tests/test_plugin.py
@@ -6,7 +6,7 @@
 from kedro import __version__ as kedro_version
 from kedro.framework.project import pipelines
 from kedro.framework.startup import ProjectMetadata
-from kedro.io import DataCatalog, MemoryDataSet
+from kedro.io import DataCatalog, MemoryDataset
 from kedro.pipeline import node
 from kedro.pipeline.modular_pipeline import pipeline as modular_pipeline
 from pytest import fixture
@@ -40,9 +40,9 @@ def fake_metadata(tmp_path):
 @fixture
 def fake_context(mocker):
     mock_context = mocker.Mock()
-    dummy_1 = MemoryDataSet()
-    dummy_2 = MemoryDataSet()
-    dummy_3 = MemoryDataSet()
+    dummy_1 = MemoryDataset()
+    dummy_2 = MemoryDataset()
+    dummy_3 = MemoryDataset()
     mock_context.catalog = DataCatalog(
         {"dummy_1": dummy_1, "dummy_2": dummy_2, "dummy_3": dummy_3}
     )