Function TEST_F

tensorflow/core/grappler/grappler_item_builder_test.cc:38–95

class GrapplerItemBuilderTest : public ::testing::Test {};

TEST_F(GrapplerItemBuilderTest, AssetFilepathOverrideTest) {
  MetaGraphDef meta_graph;

  tensorflow::Scope s = tensorflow::Scope::NewRootScope();
  Output var =
      ops::Variable(s.WithOpName("var"), TensorShape(), DataType::DT_FLOAT);
  Output filename_node =
      ops::Const(s.WithOpName("filename"), string("model"), TensorShape());
  Output tensor_name =
      ops::Const(s.WithOpName("tensorname"), string("var"), TensorShape());
  Output restore = ops::Restore(s.WithOpName("restore"), filename_node,
                                tensor_name, DataType::DT_FLOAT);
  Output assign = ops::Assign(s.WithOpName("assign"), var, restore);

  TF_CHECK_OK(s.ToGraphDef(meta_graph.mutable_graph_def()));

  string temp_dir = testing::TmpDir();

  Env *env = Env::Default();
  string filename =
      io::JoinPath(temp_dir, "grappler_item_builder_test_filename");
  env->DeleteFile(filename).IgnoreError();
  std::unique_ptr<WritableFile> file_to_write;
  TF_CHECK_OK(env->NewWritableFile(filename, &file_to_write));
  TF_CHECK_OK(file_to_write->Close());
  TF_CHECK_OK(env->FileExists(filename));
  LOG(INFO) << filename;

  AssetFileDef asset_file_def;
  *asset_file_def.mutable_tensor_info()->mutable_name() = "filename";
  *asset_file_def.mutable_filename() = "grappler_item_builder_test_filename";

  (*meta_graph.mutable_collection_def())["saved_model_assets"]
      .mutable_any_list()
      ->add_value()
      ->PackFrom(asset_file_def);
  *((*meta_graph.mutable_collection_def())["train_op"]
        .mutable_node_list()
        ->add_value()) = "assign";

  ItemConfig cfg;
  cfg.assets_directory_override = temp_dir;

  std::unique_ptr<GrapplerItem> item =
      GrapplerItemFromMetaGraphDef("0", meta_graph, cfg);
  ASSERT_TRUE(item != nullptr);
  for (const NodeDef &node : item->graph.node()) {
    if (node.name() == "filename") {
      const auto iter = node.attr().find("value");
      ASSERT_TRUE(iter != node.attr().end());
      ASSERT_TRUE(iter->second.has_tensor());
      ASSERT_EQ(1, iter->second.tensor().string_val_size());

      string tensor_string_val = iter->second.tensor().string_val(0);
      EXPECT_EQ(tensor_string_val, filename);
    }
  }
}