Making sure that the src_incarnation field on the ParsedKey for the Send and Recv's is set correctly.

PiperOrigin-RevId: 168635306

1 files changed, 14 insertions, 0 deletions

diff --git a/tensorflow/core/common_runtime/process_function_library_runtime_test.cc b/tensorflow/core/common_runtime/process_function_library_runtime_test.cc
index fdbab46f54..50379a52c4 100644
--- a/tensorflow/core/common_runtime/process_function_library_runtime_test.cc
+++ b/tensorflow/core/common_runtime/process_function_library_runtime_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/rendezvous_mgr.h"
 #include "tensorflow/core/framework/function_testlib.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/public/session_options.h"
 #include "tensorflow/core/public/version.h"
@@ -120,6 +121,19 @@ TEST_F(ProcessFunctionLibraryRuntimeTest, ObtainFunctionTarget) {
   EXPECT_EQ("/job:a/replica:0/task:0/cpu:1", target);
 }
 
+TEST_F(ProcessFunctionLibraryRuntimeTest, GetDeviceIncarnation) {
+  Init({});
+  int64 incarnation;
+  TF_EXPECT_OK(proc_flr_->GetDeviceIncarnation("/job:a/replica:0/task:0/cpu:1",
+                                               &incarnation));
+  // Incarnation is a random number other than 0.
+  EXPECT_NE(incarnation, 0);
+  Status s = proc_flr_->GetDeviceIncarnation("/job:a/replica:0/task:0/cpu:2",
+                                             &incarnation);
+  EXPECT_EQ(s.code(), error::INVALID_ARGUMENT);
+  rendezvous_->Unref();
+}
+
 TEST_F(ProcessFunctionLibraryRuntimeTest, SingleCall) {
   Init({test::function::XTimesTwo()});
   FunctionLibraryRuntime::Options opts;