diff --git a/mindspore/ccsrc/minddata/dataset/engine/datasetops/dataset_op.cc b/mindspore/ccsrc/minddata/dataset/engine/datasetops/dataset_op.cc
index cda64420eb..951d28c4f4 100644
--- a/mindspore/ccsrc/minddata/dataset/engine/datasetops/dataset_op.cc
+++ b/mindspore/ccsrc/minddata/dataset/engine/datasetops/dataset_op.cc
@@ -108,7 +108,7 @@ Status DatasetOp::InsertAsParent(std::shared_ptr<DatasetOp> to_add) {
   }
   RETURN_IF_NOT_OK(to_add->AddChild(shared_from_this()));
   if (tree_->root()->id() == this->id()) {
-    tree_->AssignRoot(to_add);
+    RETURN_IF_NOT_OK(tree_->AssignRoot(to_add));
   }
   return Status::OK();
 }
@@ -156,7 +156,7 @@ Status DatasetOp::Remove() {
       // We don't have a parent, so we are the root node being removed.
       // clear the parent list of our child so that it becomes the new root.
       child_[0]->parent_.clear();
-      tree_->AssignRoot(child_[0]);
+      RETURN_IF_NOT_OK(tree_->AssignRoot(child_[0]));
     }
   }
 
diff --git a/mindspore/ccsrc/minddata/dataset/engine/opt/pre/epoch_injection_pass.cc b/mindspore/ccsrc/minddata/dataset/engine/opt/pre/epoch_injection_pass.cc
index 34e225de62..17f15f8f81 100644
--- a/mindspore/ccsrc/minddata/dataset/engine/opt/pre/epoch_injection_pass.cc
+++ b/mindspore/ccsrc/minddata/dataset/engine/opt/pre/epoch_injection_pass.cc
@@ -68,7 +68,7 @@ Status EpochInjectionPass::RunOnTree(ExecutionTree *tree, bool *modified) {
     std::shared_ptr<EpochCtrlOp> epoch_ctrl_op;
     RETURN_IF_NOT_OK(EpochCtrlOp::Builder(num_epochs).Build(&epoch_ctrl_op));
     RETURN_IF_NOT_OK(tree->AssociateNode(epoch_ctrl_op));
-    epoch_inject_node->InsertAsParent(epoch_ctrl_op);
+    RETURN_IF_NOT_OK(epoch_inject_node->InsertAsParent(epoch_ctrl_op));
   }
 
   MS_LOG(INFO) << "Pre pass: Injection pass complete.";
diff --git a/tests/ut/cpp/dataset/album_op_test.cc b/tests/ut/cpp/dataset/album_op_test.cc
index dadfe7d9a6..49af1a63d0 100644
--- a/tests/ut/cpp/dataset/album_op_test.cc
+++ b/tests/ut/cpp/dataset/album_op_test.cc
@@ -90,56 +90,44 @@ TEST_F(MindDataTestAlbum, TestSequentialAlbumWithSchema) {
   std::string schema_file = datasets_root_path_ + "/testAlbum/datasetSchema.json";
   std::vector<std::string> column_names = {"image", "label", "id"};
   auto tree = Build({AlbumSchema(16, 2, 32, folder_path, schema_file, column_names, false), Repeat(2)});
-  tree->Prepare();
-  Status rc = tree->Launch();
-  if (rc.IsError()) {
-    MS_LOG(ERROR) << "Return code error detected during tree launch: " <<  ".";
-    EXPECT_TRUE(false);
-  } else {
-    DatasetIterator di(tree);
-    TensorMap tensor_map;
+  ASSERT_OK(tree->Prepare());
+  ASSERT_OK(tree->Launch());
+  DatasetIterator di(tree);
+  TensorMap tensor_map;
+  ASSERT_OK(di.GetNextAsMap(&tensor_map));
+  uint64_t i = 0;
+  std::string_view label = 0;
+  while (tensor_map.size() != 0) {
+    EXPECT_TRUE(tensor_map["label"]->GetItemAt(&label, {}));
+    MS_LOG(DEBUG) << "row: " << i << "\t" << tensor_map["image"]->shape() << "label:" << label << "label shape"
+                  << tensor_map["label"] << "\n";
+    i++;
     di.GetNextAsMap(&tensor_map);
-    EXPECT_TRUE(rc.IsOk());
-    uint64_t i = 0;
-    int32_t label = 0;
-    while (tensor_map.size() != 0) {
-      tensor_map["label"]->GetItemAt<int32_t>(&label, {});
-      MS_LOG(DEBUG) << "row: " << i << "\t" << tensor_map["image"]->shape() << "label:" << label << "label shape"
-                    << tensor_map["label"] << "\n";
-      i++;
-      di.GetNextAsMap(&tensor_map);
-    }
-    MS_LOG(INFO) << "got rows" << i << "\n";
-    EXPECT_TRUE(i == 14);
   }
+  MS_LOG(INFO) << "got rows: " << i << "\n";
+  EXPECT_TRUE(i == 14);
 }
 
 TEST_F(MindDataTestAlbum, TestSequentialAlbumWithSchemaNoOrder) {
   std::string folder_path = datasets_root_path_ + "/testAlbum/images";
   std::string schema_file = datasets_root_path_ + "/testAlbum/datasetSchema.json";
   auto tree = Build({AlbumSchema(16, 2, 32, folder_path, schema_file), Repeat(2)});
-  tree->Prepare();
-  Status rc = tree->Launch();
-  if (rc.IsError()) {
-    MS_LOG(ERROR) << "Return code error detected during tree launch: " << ".";
-    EXPECT_TRUE(false);
-  } else {
-    DatasetIterator di(tree);
-    TensorMap tensor_map;
+  ASSERT_OK(tree->Prepare());
+  ASSERT_OK(tree->Launch());
+  DatasetIterator di(tree);
+  TensorMap tensor_map;
+  ASSERT_TRUE(di.GetNextAsMap(&tensor_map));
+  uint64_t i = 0;
+  std::string_view label;
+  while (tensor_map.size() != 0) {
+    EXPECT_OK(tensor_map["label"]->GetItemAt(&label, {}));
+    MS_LOG(DEBUG) << "row: " << i << "\t" << tensor_map["image"]->shape() << "label:" << label << "label shape"
+                  << tensor_map["label"] << "\n";
+    i++;
     di.GetNextAsMap(&tensor_map);
-    EXPECT_TRUE(rc.IsOk());
-    uint64_t i = 0;
-    int32_t label = 0;
-    while (tensor_map.size() != 0) {
-      tensor_map["label"]->GetItemAt<int32_t>(&label, {});
-      MS_LOG(DEBUG) << "row: " << i << "\t" << tensor_map["image"]->shape() << "label:" << label << "label shape"
-                    << tensor_map["label"] << "\n";
-      i++;
-      di.GetNextAsMap(&tensor_map);
-    }
-    MS_LOG(INFO) << "got rows" << i << "\n";
-    EXPECT_TRUE(i == 14);
   }
+  MS_LOG(INFO) << "got rows: " << i << "\n";
+  EXPECT_TRUE(i == 14);
 }
 
 TEST_F(MindDataTestAlbum, TestSequentialAlbumWithSchemaFloat) {
@@ -148,29 +136,23 @@ TEST_F(MindDataTestAlbum, TestSequentialAlbumWithSchemaFloat) {
   std::string schema_file = datasets_root_path_ + "/testAlbum/floatSchema.json";
   auto tree = Build({AlbumSchema(16, 2, 32, folder_path, schema_file), Repeat(2)});
   tree->Prepare();
-  Status rc = tree->Launch();
-  if (rc.IsError()) {
-    MS_LOG(ERROR) << "Return code error detected during tree launch: " << ".";
-    EXPECT_TRUE(false);
-  } else {
-    DatasetIterator di(tree);
-    TensorMap tensor_map;
+  ASSERT_OK(tree->Launch());
+  DatasetIterator di(tree);
+  TensorMap tensor_map;
+  ASSERT_OK(di.GetNextAsMap(&tensor_map));
+  uint64_t i = 0;
+  std::string_view label;
+  double priority = 0;
+  while (tensor_map.size() != 0) {
+    EXPECT_OK(tensor_map["label"]->GetItemAt(&label, {}));
+    EXPECT_OK(tensor_map["_priority"]->GetItemAt<double>(&priority, {}));
+    MS_LOG(DEBUG) << "row: " << i << "\t" << tensor_map["image"]->shape() << "label:" << label << "label shape"
+                  << tensor_map["label"]  << "priority: " << priority << "\n";
+    i++;
     di.GetNextAsMap(&tensor_map);
-    EXPECT_TRUE(rc.IsOk());
-    uint64_t i = 0;
-    int32_t label = 0;
-    double priority = 0;
-    while (tensor_map.size() != 0) {
-      tensor_map["label"]->GetItemAt<int32_t>(&label, {});
-      tensor_map["_priority"]->GetItemAt<double>(&priority, {});
-      MS_LOG(DEBUG) << "row: " << i << "\t" << tensor_map["image"]->shape() << "label:" << label << "label shape"
-                    << tensor_map["label"]  << "priority: " << priority << "\n";
-      i++;
-      di.GetNextAsMap(&tensor_map);
-    }
-    MS_LOG(INFO) << "got rows" << i << "\n";
-    EXPECT_TRUE(i == 14);
   }
+  MS_LOG(INFO) << "got rows: " << i << "\n";
+  EXPECT_TRUE(i == 14);
 }
 
 TEST_F(MindDataTestAlbum, TestSequentialAlbumWithFullSchema) {
@@ -178,32 +160,26 @@ TEST_F(MindDataTestAlbum, TestSequentialAlbumWithFullSchema) {
   // add the priority column
   std::string schema_file = datasets_root_path_ + "/testAlbum/fullSchema.json";
   auto tree = Build({AlbumSchema(16, 2, 32, folder_path, schema_file), Repeat(2)});
-  tree->Prepare();
-  Status rc = tree->Launch();
-  if (rc.IsError()) {
-    MS_LOG(ERROR) << "Return code error detected during tree launch: " << ".";
-    EXPECT_TRUE(false);
-  } else {
-    DatasetIterator di(tree);
-    TensorMap tensor_map;
+  ASSERT_OK(tree->Prepare());
+  ASSERT_OK(tree->Launch());
+  DatasetIterator di(tree);
+  TensorMap tensor_map;
+  ASSERT_OK(di.GetNextAsMap(&tensor_map));
+  uint64_t i = 0;
+  std::string_view label = 0;
+  double priority = 0;
+  int64_t id = 0;
+  while (tensor_map.size() != 0) {
+    EXPECT_OK(tensor_map["label"]->GetItemAt(&label, {}));
+    EXPECT_OK(tensor_map["_priority"]->GetItemAt<double>(&priority, {}));
+    EXPECT_OK(tensor_map["id"]->GetItemAt<int64_t>(&id, {}));
+    MS_LOG(DEBUG) << "row: " << i << "\t" << tensor_map["image"]->shape() << "label:" << label << "label shape"
+                  << tensor_map["label"]  << "priority: " << priority << " embedding : " 
+                  << tensor_map["_embedding"]->shape() << " id: " << id << "\n";
+    i++;
     di.GetNextAsMap(&tensor_map);
-    EXPECT_TRUE(rc.IsOk());
-    uint64_t i = 0;
-    int32_t label = 0;
-    double priority = 0;
-    int64_t id = 0;
-    while (tensor_map.size() != 0) {
-      tensor_map["label"]->GetItemAt<int32_t>(&label, {});
-      tensor_map["_priority"]->GetItemAt<double>(&priority, {});
-      tensor_map["id"]->GetItemAt<int64_t>(&id, {});
-      MS_LOG(ERROR) << "row: " << i << "\t" << tensor_map["image"]->shape() << "label:" << label << "label shape"
-                    << tensor_map["label"]  << "priority: " << priority << " embedding : " <<
-		    tensor_map["_embedding"]->shape() << " id: " << id << "\n";
-      i++;
-      di.GetNextAsMap(&tensor_map);
-    }
-    MS_LOG(INFO) << "got rows" << i << "\n";
-    EXPECT_TRUE(i == 14);
   }
+  MS_LOG(INFO) << "got rows: " << i << "\n";
+  EXPECT_TRUE(i == 14);
 }
 
diff --git a/tests/ut/cpp/dataset/common/common.h b/tests/ut/cpp/dataset/common/common.h
index 539978177a..dc865cab05 100644
--- a/tests/ut/cpp/dataset/common/common.h
+++ b/tests/ut/cpp/dataset/common/common.h
@@ -17,8 +17,27 @@
 #define TESTS_DATASET_UT_CORE_COMMON_DE_UT_COMMON_H_
 
 #include "gtest/gtest.h"
+#include "minddata/dataset/util/status.h"
 #include "utils/log_adapter.h"
 
+#define ASSERT_OK(_s)                          \
+  do {                                         \
+    Status __rc = (_s);                        \
+    if (__rc.IsError()) {                      \
+      MS_LOG(ERROR) << __rc.ToString() << "."; \
+      ASSERT_TRUE(false);                      \
+    }                                          \
+  } while (false)
+
+#define EXPECT_OK(_s)                          \
+  do {                                         \
+    Status __rc = (_s);                        \
+    if (__rc.IsError()) {                      \
+      MS_LOG(ERROR) << __rc.ToString() << "."; \
+      EXPECT_TRUE(false);                      \
+    }                                          \
+  } while (false)
+
 namespace UT {
 class Common : public testing::Test {
  public:
diff --git a/tests/ut/cpp/dataset/data_helper_test.cc b/tests/ut/cpp/dataset/data_helper_test.cc
index 3262279b66..ea8c79c633 100644
--- a/tests/ut/cpp/dataset/data_helper_test.cc
+++ b/tests/ut/cpp/dataset/data_helper_test.cc
@@ -153,6 +153,10 @@ TEST_F(MindDataTestDataHelper, MindDataTestTensorWriteFloat) {
   // create buffer using system mempool
   DataHelper dh;
   void *data = malloc(t->SizeInBytes());
+  if (data == nullptr) {
+    MS_LOG(ERROR) << "malloc failed";
+    ASSERT_TRUE(false);
+  }
   auto bytes_copied = dh.DumpData(t->GetBuffer(), t->SizeInBytes(), data, t->SizeInBytes());
   if (bytes_copied != t->SizeInBytes()) {
     EXPECT_TRUE(false);
@@ -177,6 +181,10 @@ TEST_F(MindDataTestDataHelper, MindDataTestTensorWriteUInt) {
   // create buffer using system mempool
   DataHelper dh;
   void *data = malloc(t->SizeInBytes());
+  if (data == nullptr) {
+    MS_LOG(ERROR) << "malloc failed";
+    ASSERT_TRUE(false);
+  }
   auto bytes_copied = dh.DumpData(t->GetBuffer(), t->SizeInBytes(), data, t->SizeInBytes());
   if (bytes_copied != t->SizeInBytes()) {
     EXPECT_TRUE(false);