xref: /aosp_15_r20/external/federated-compute/fcp/tensorflow/host_object_test.cc (revision 14675a029014e728ec732f129a32e299b2da0601)
1 /*
2  * Copyright 2019 Google LLC
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #include "fcp/tensorflow/host_object.h"
18 
19 #include "gmock/gmock.h"
20 #include "gtest/gtest.h"
21 
22 namespace fcp {
23 
24 using ::testing::Eq;
25 
26 class WidgetInterface {
27  public:
28   virtual ~WidgetInterface() = default;
29   virtual void Poke(int value) = 0;
30 };
31 
32 class WidgetImpl : public WidgetInterface {
33  public:
Poke(int value)34   void Poke(int value) final {
35     counter_ += value;
36   }
37 
counter() const38   int counter() const {
39     return counter_;
40   }
41  private:
42   int counter_ = 0;
43 };
44 
TEST(HostObjectTest,LookupFailure)45 TEST(HostObjectTest, LookupFailure) {
46   std::optional<std::shared_ptr<WidgetInterface>> p =
47       HostObjectRegistry<WidgetInterface>::TryLookup(RandomToken::Generate());
48   EXPECT_THAT(p, Eq(std::nullopt));
49 }
50 
TEST(HostObjectTest,LookupSuccess)51 TEST(HostObjectTest, LookupSuccess) {
52   std::shared_ptr<WidgetImpl> obj = std::make_shared<WidgetImpl>();
53   HostObjectRegistration reg =
54       HostObjectRegistry<WidgetInterface>::Register(obj);
55 
56   std::optional<std::shared_ptr<WidgetInterface>> p =
57       HostObjectRegistry<WidgetInterface>::TryLookup(reg.token());
58   EXPECT_TRUE(p.has_value());
59 
60   (*p)->Poke(123);
61   EXPECT_THAT(obj->counter(), Eq(123));
62   EXPECT_THAT(p->get(), Eq(obj.get()));
63 }
64 
TEST(HostObjectTest,Unregister)65 TEST(HostObjectTest, Unregister) {
66   std::shared_ptr<WidgetImpl> obj = std::make_shared<WidgetImpl>();
67 
68   std::optional<RandomToken> token;
69   {
70     HostObjectRegistration reg =
71         HostObjectRegistry<WidgetInterface>::Register(obj);
72     token = reg.token();
73   }
74 
75   std::optional<std::shared_ptr<WidgetInterface>> p =
76       HostObjectRegistry<WidgetInterface>::TryLookup(*token);
77   EXPECT_THAT(p, Eq(std::nullopt));
78 }
79 
80 }  // namespace fcp
81