| 80 | } // namespace |
| 81 | |
| 82 | TEST(SetDeviceTest, TestSetDevice) { |
| 83 | GraphDef graph_def = CreateDeviceGraph(); |
| 84 | GraphDef result; |
| 85 | TransformFuncContext context; |
| 86 | context.input_names = {}; |
| 87 | context.output_names = {"mul_node1"}; |
| 88 | context.params.insert(std::pair<string, std::vector<string>>( |
| 89 | {"device", {string("/device:CPU:0")}})); |
| 90 | TF_ASSERT_OK(SetDevice(graph_def, context, &result)); |
| 91 | |
| 92 | std::map<string, const NodeDef*> node_lookup; |
| 93 | MapNamesToNodes(result, &node_lookup); |
| 94 | EXPECT_EQ("/device:CPU:0", node_lookup.at("mul_node1")->device()); |
| 95 | EXPECT_EQ("/device:CPU:0", node_lookup.at("add_node2")->device()); |
| 96 | EXPECT_EQ("/device:CPU:0", node_lookup.at("add_node3")->device()); |
| 97 | EXPECT_EQ("/device:CPU:0", node_lookup.at("const_node1")->device()); |
| 98 | EXPECT_EQ("/device:CPU:0", node_lookup.at("const_node2")->device()); |
| 99 | EXPECT_EQ("/device:CPU:0", node_lookup.at("const_node3")->device()); |
| 100 | EXPECT_EQ("/device:CPU:0", node_lookup.at("add_node4")->device()); |
| 101 | } |
| 102 | |
| 103 | TEST(SetDeviceTest, TestSetDeviceIfDefault) { |
| 104 | GraphDef graph_def = CreateDeviceGraph(); |
nothing calls this directly
no test coverage detected