| 30 | namespace { |
| 31 | |
| 32 | class NodeDefBuilderTest : public ::testing::Test { |
| 33 | protected: |
| 34 | // Specify an OpDef via an OpDefBuilder. |
| 35 | void Op(const OpDefBuilder& op_def_builder) { |
| 36 | OpRegistrationData op_reg_data; |
| 37 | TF_EXPECT_OK(op_def_builder.Finalize(&op_reg_data)); |
| 38 | op_def_ = op_reg_data.op_def; |
| 39 | } |
| 40 | |
| 41 | // Resets builder_ with a new NodeDefBuilder using the Op from the last call |
| 42 | // to Op() above. |
| 43 | NodeDefBuilder& Builder() { |
| 44 | EXPECT_FALSE(op_def_.name().empty()) << "Must call Op() before Builder()"; |
| 45 | builder_.reset(new NodeDefBuilder("n", &op_def_)); |
| 46 | return *builder_; |
| 47 | } |
| 48 | |
| 49 | // Calls Finalize() and verifies it returns success and the result matches |
| 50 | // expectations. |
| 51 | void ExpectSuccess(NodeDefBuilder& builder, // NOLINT |
| 52 | DataTypeSlice expected_in_types, |
| 53 | DataTypeSlice expected_out_types, StringPiece proto) { |
| 54 | NodeDef node_def; |
| 55 | Status status = builder.Finalize(&node_def); |
| 56 | TF_EXPECT_OK(status); |
| 57 | if (!status.ok()) return; |
| 58 | NodeDef expected; |
| 59 | protobuf::TextFormat::ParseFromString(strings::StrCat("name: 'n' ", proto), |
| 60 | &expected); |
| 61 | EXPECT_EQ(node_def.DebugString(), expected.DebugString()); |
| 62 | |
| 63 | DataTypeVector in_types, out_types; |
| 64 | status = |
| 65 | InOutTypesForNode(node_def, builder.op_def(), &in_types, &out_types); |
| 66 | TF_EXPECT_OK(status); |
| 67 | if (!status.ok()) return; |
| 68 | EXPECT_EQ(DataTypeSliceString(expected_in_types), |
| 69 | DataTypeVectorString(in_types)); |
| 70 | EXPECT_EQ(DataTypeSliceString(expected_out_types), |
| 71 | DataTypeVectorString(out_types)); |
| 72 | |
| 73 | status = ValidateNodeDef(node_def, op_def_); |
| 74 | TF_EXPECT_OK(status); |
| 75 | } |
| 76 | |
| 77 | // Calls Finalize() and verifies it returns an error. |
| 78 | // Each message must appear as a substring of the error. |
| 79 | void ExpectFailures(NodeDefBuilder& builder, // NOLINT |
| 80 | const std::vector<string>& messages) { |
| 81 | NodeDef node_def; |
| 82 | Status status = builder.Finalize(&node_def); |
| 83 | EXPECT_FALSE(status.ok()) << SummarizeNodeDef(node_def); |
| 84 | if (status.ok()) return; |
| 85 | for (const string& message : messages) { |
| 86 | EXPECT_TRUE(absl::StrContains(status.error_message(), message)) |
| 87 | << status << ", " << message; |
| 88 | } |
| 89 | } |