diff --git a/libs/binder/tests/binderLibTest.cpp b/libs/binder/tests/binderLibTest.cpp
index 7bcfffd..22c1bad 100644
--- a/libs/binder/tests/binderLibTest.cpp
+++ b/libs/binder/tests/binderLibTest.cpp
@@ -953,7 +953,7 @@
 {
     status_t ret;
     Parcel data, reply;
-    uid_t previousWorkSource = IPCThreadState::self()->setWorkSource(100);
+    int64_t previousWorkSource = IPCThreadState::self()->setCallingWorkSourceUid(100);
     data.writeInterfaceToken(binderLibTestServiceName);
     ret = m_server->transact(BINDER_LIB_TEST_GET_WORK_SOURCE_TRANSACTION, data, &reply);
     EXPECT_EQ(100, reply.readInt32());
@@ -966,8 +966,8 @@
     status_t ret;
     Parcel data, reply;
 
-    IPCThreadState::self()->setWorkSource(100);
-    uid_t previousWorkSource = IPCThreadState::self()->clearWorkSource();
+    IPCThreadState::self()->setCallingWorkSourceUid(100);
+    int64_t previousWorkSource = IPCThreadState::self()->clearCallingWorkSource();
     data.writeInterfaceToken(binderLibTestServiceName);
     ret = m_server->transact(BINDER_LIB_TEST_GET_WORK_SOURCE_TRANSACTION, data, &reply);
 
@@ -976,6 +976,22 @@
     EXPECT_EQ(NO_ERROR, ret);
 }
 
+TEST_F(BinderLibTest, WorkSourceRestored)
+{
+    status_t ret;
+    Parcel data, reply;
+
+    IPCThreadState::self()->setCallingWorkSourceUid(100);
+    int64_t token = IPCThreadState::self()->clearCallingWorkSource();
+    IPCThreadState::self()->restoreCallingWorkSource(token);
+
+    data.writeInterfaceToken(binderLibTestServiceName);
+    ret = m_server->transact(BINDER_LIB_TEST_GET_WORK_SOURCE_TRANSACTION, data, &reply);
+
+    EXPECT_EQ(100, reply.readInt32());
+    EXPECT_EQ(NO_ERROR, ret);
+}
+
 class BinderLibTestService : public BBinder
 {
     public:
@@ -1276,7 +1292,7 @@
             }
             case BINDER_LIB_TEST_GET_WORK_SOURCE_TRANSACTION: {
                 data.enforceInterface(binderLibTestServiceName);
-                reply->writeInt32(IPCThreadState::self()->getWorkSource());
+                reply->writeInt32(IPCThreadState::self()->getCallingWorkSourceUid());
                 return NO_ERROR;
             }
             default: