Add support of Components/Traits to local privet APIs

Added support for new privet APIs /privet/v3/traits and
/privet/v3/components to obtain the full trait/component trees
as well as expanded the existing /privet/v3/checkForUpdates to
include the component/trait fingerprints.

BUG: 25917521
Change-Id: Ib753817f88f611935057ca7bd1a1bf02addfb69c
Reviewed-on: https://weave-review.googlesource.com/1791
Reviewed-by: Vitaly Buka <vitalybuka@google.com>

6 files changed, 296 insertions, 61 deletions

diff --git a/src/privet/cloud_delegate.cc b/src/privet/cloud_delegate.cc
index 9c8a619..170688b 100644
--- a/src/privet/cloud_delegate.cc
+++ b/src/privet/cloud_delegate.cc
@@ -58,7 +58,7 @@ class CloudDelegateImpl : public CloudDelegate {
     component_manager_->AddCommandRemovedCallback(base::Bind(
         &CloudDelegateImpl::OnCommandRemoved, weak_factory_.GetWeakPtr()));
     component_manager_->AddStateChangedCallback(base::Bind(
-        &CloudDelegateImpl::NotifyOnComponentTreeChanged,
+        &CloudDelegateImpl::NotifyOnStateChanged,
         weak_factory_.GetWeakPtr()));
     component_manager_->AddComponentTreeChangedCallback(base::Bind(
         &CloudDelegateImpl::NotifyOnComponentTreeChanged,
@@ -375,5 +375,9 @@ void CloudDelegate::NotifyOnComponentTreeChanged() {
   FOR_EACH_OBSERVER(Observer, observer_list_, OnComponentTreeChanged());
 }
 
+void CloudDelegate::NotifyOnStateChanged() {
+  FOR_EACH_OBSERVER(Observer, observer_list_, OnStateChanged());
+}
+
 }  // namespace privet
 }  // namespace weave
diff --git a/src/privet/cloud_delegate.h b/src/privet/cloud_delegate.h
index e1b1887..e80c39d 100644
--- a/src/privet/cloud_delegate.h
+++ b/src/privet/cloud_delegate.h
@@ -48,6 +48,7 @@ class CloudDelegate {
 
     virtual void OnDeviceInfoChanged() {}
     virtual void OnTraitDefsChanged() {}
+    virtual void OnStateChanged() {}
     virtual void OnComponentTreeChanged() {}
   };
 
@@ -132,6 +133,7 @@ class CloudDelegate {
 
   void NotifyOnDeviceInfoChanged();
   void NotifyOnTraitDefsChanged();
+  void NotifyOnStateChanged();
   void NotifyOnComponentTreeChanged();
 
   // Create default instance.
diff --git a/src/privet/privet_handler.cc b/src/privet/privet_handler.cc
index d609787..9199c6f 100644
--- a/src/privet/privet_handler.cc
+++ b/src/privet/privet_handler.cc
@@ -103,10 +103,14 @@ const char kSetupStartUserKey[] = "user";
 const char kFingerprintKey[] = "fingerprint";
 const char kStateKey[] = "state";
 const char kCommandsKey[] = "commands";
+const char kTraitsKey[] = "traits";
+const char kComponentsKey[] = "components";
 const char kCommandsIdKey[] = "id";
 
 const char kStateFingerprintKey[] = "stateFingerprint";
 const char kCommandsFingerprintKey[] = "commandsFingerprint";
+const char kTraitsFingerprintKey[] = "traitsFingerprint";
+const char kComponentsFingerprintKey[] = "componentsFingerprint";
 const char kWaitTimeoutKey[] = "waitTimeout";
 
 const char kInvalidParamValueFormat[] = "Invalid parameter: '%s'='%s'";
@@ -371,6 +375,10 @@ PrivetHandler::PrivetHandler(CloudDelegate* cloud,
                    &PrivetHandler::HandleCommandsList, AuthScope::kViewer);
   AddSecureHandler("/privet/v3/checkForUpdates",
                    &PrivetHandler::HandleCheckForUpdates, AuthScope::kViewer);
+  AddSecureHandler("/privet/v3/traits", &PrivetHandler::HandleTraits,
+                   AuthScope::kViewer);
+  AddSecureHandler("/privet/v3/components", &PrivetHandler::HandleComponents,
+                   AuthScope::kViewer);
 }
 
 PrivetHandler::~PrivetHandler() {
@@ -379,9 +387,9 @@ PrivetHandler::~PrivetHandler() {
 }
 
 void PrivetHandler::OnTraitDefsChanged() {
-  ++command_defs_fingerprint_;
+  ++traits_fingerprint_;
   auto pred = [this](const UpdateRequestParameters& params) {
-    return params.command_defs_fingerprint < 0;
+    return params.traits_fingerprint < 0;
   };
   auto last =
       std::partition(update_requests_.begin(), update_requests_.end(), pred);
@@ -390,10 +398,24 @@ void PrivetHandler::OnTraitDefsChanged() {
   update_requests_.erase(last, update_requests_.end());
 }
 
-void PrivetHandler::OnComponentTreeChanged() {
+void PrivetHandler::OnStateChanged() {
+  // State updates also change the component tree, so update both fingerprints.
   ++state_fingerprint_;
+  ++components_fingerprint_;
+  auto pred = [this](const UpdateRequestParameters& params) {
+    return params.state_fingerprint < 0 && params.components_fingerprint < 0;
+  };
+  auto last =
+      std::partition(update_requests_.begin(), update_requests_.end(), pred);
+  for (auto p = last; p != update_requests_.end(); ++p)
+    ReplyToUpdateRequest(p->callback);
+  update_requests_.erase(last, update_requests_.end());
+}
+
+void PrivetHandler::OnComponentTreeChanged() {
+  ++components_fingerprint_;
   auto pred = [this](const UpdateRequestParameters& params) {
-    return params.state_fingerprint < 0;
+    return params.components_fingerprint < 0;
   };
   auto last =
       std::partition(update_requests_.begin(), update_requests_.end(), pred);
@@ -768,12 +790,34 @@ void PrivetHandler::HandleState(const base::DictionaryValue& input,
   callback.Run(http::kOk, output);
 }
 
+void PrivetHandler::HandleTraits(const base::DictionaryValue& input,
+                                 const UserInfo& user_info,
+                                 const RequestCallback& callback) {
+  base::DictionaryValue output;
+  output.Set(kTraitsKey, cloud_->GetTraits().DeepCopy());
+  output.SetString(kFingerprintKey, std::to_string(traits_fingerprint_));
+
+  callback.Run(http::kOk, output);
+}
+
+void PrivetHandler::HandleComponents(const base::DictionaryValue& input,
+                                     const UserInfo& user_info,
+                                     const RequestCallback& callback) {
+  base::DictionaryValue output;
+  output.Set(kComponentsKey, cloud_->GetComponents().DeepCopy());
+  output.SetString(kFingerprintKey, std::to_string(components_fingerprint_));
+
+  callback.Run(http::kOk, output);
+}
+
 void PrivetHandler::HandleCommandDefs(const base::DictionaryValue& input,
                                       const UserInfo& user_info,
                                       const RequestCallback& callback) {
   base::DictionaryValue output;
   output.Set(kCommandsKey, cloud_->GetLegacyCommandDef().DeepCopy());
-  output.SetString(kFingerprintKey, std::to_string(command_defs_fingerprint_));
+  // Use traits fingerprint since right now we treat traits and command defs
+  // as being equivalent.
+  output.SetString(kFingerprintKey, std::to_string(traits_fingerprint_));
 
   callback.Run(http::kOk, output);
 }
@@ -845,12 +889,18 @@ void PrivetHandler::HandleCheckForUpdates(const base::DictionaryValue& input,
 
   std::string state_fingerprint;
   std::string commands_fingerprint;
+  std::string traits_fingerprint;
+  std::string components_fingerprint;
   input.GetString(kStateFingerprintKey, &state_fingerprint);
   input.GetString(kCommandsFingerprintKey, &commands_fingerprint);
+  input.GetString(kTraitsFingerprintKey, &traits_fingerprint);
+  input.GetString(kComponentsFingerprintKey, &components_fingerprint);
   const bool ignore_state = state_fingerprint.empty();
   const bool ignore_commands = commands_fingerprint.empty();
-  // If both fingerprints are missing, nothing to wait for, return immediately.
-  if (ignore_state && ignore_commands)
+  const bool ignore_traits = traits_fingerprint.empty();
+  const bool ignore_components = components_fingerprint.empty();
+  // If all fingerprints are missing, nothing to wait for, return immediately.
+  if (ignore_state && ignore_commands && ignore_traits && ignore_components)
     return ReplyToUpdateRequest(callback);
   // If the current state fingerprint is different from the requested one,
   // return new fingerprints.
@@ -858,17 +908,32 @@ void PrivetHandler::HandleCheckForUpdates(const base::DictionaryValue& input,
     return ReplyToUpdateRequest(callback);
   // If the current commands fingerprint is different from the requested one,
   // return new fingerprints.
+  // NOTE: We are using traits fingerprint for command fingerprint as well.
   if (!ignore_commands &&
-      commands_fingerprint != std::to_string(command_defs_fingerprint_)) {
+      commands_fingerprint != std::to_string(traits_fingerprint_)) {
+    return ReplyToUpdateRequest(callback);
+  }
+  // If the current traits fingerprint is different from the requested one,
+  // return new fingerprints.
+  if (!ignore_traits &&
+      traits_fingerprint != std::to_string(traits_fingerprint_)) {
+    return ReplyToUpdateRequest(callback);
+  }
+  // If the current components fingerprint is different from the requested one,
+  // return new fingerprints.
+  if (!ignore_components &&
+      components_fingerprint != std::to_string(components_fingerprint_)) {
     return ReplyToUpdateRequest(callback);
   }
 
   UpdateRequestParameters params;
   params.request_id = ++last_update_request_id_;
   params.callback = callback;
-  params.command_defs_fingerprint =
-      ignore_commands ? -1 : command_defs_fingerprint_;
+  params.traits_fingerprint =
+      (ignore_traits && ignore_commands) ? -1 : traits_fingerprint_;
   params.state_fingerprint = ignore_state ? -1 : state_fingerprint_;
+  params.components_fingerprint =
+      ignore_components ? -1 : components_fingerprint_;
   update_requests_.push_back(params);
   if (timeout != base::TimeDelta::Max()) {
     device_->PostDelayedTask(
@@ -884,7 +949,11 @@ void PrivetHandler::ReplyToUpdateRequest(
   base::DictionaryValue output;
   output.SetString(kStateFingerprintKey, std::to_string(state_fingerprint_));
   output.SetString(kCommandsFingerprintKey,
-                   std::to_string(command_defs_fingerprint_));
+                   std::to_string(traits_fingerprint_));
+  output.SetString(kTraitsFingerprintKey,
+                   std::to_string(traits_fingerprint_));
+  output.SetString(kComponentsFingerprintKey,
+                   std::to_string(components_fingerprint_));
   callback.Run(http::kOk, output);
 }
 
diff --git a/src/privet/privet_handler.h b/src/privet/privet_handler.h
index cc80d43..6bc4ac6 100644
--- a/src/privet/privet_handler.h
+++ b/src/privet/privet_handler.h
@@ -46,6 +46,7 @@ class PrivetHandler : public CloudDelegate::Observer {
   ~PrivetHandler() override;
 
   void OnTraitDefsChanged() override;
+  void OnStateChanged() override;
   void OnComponentTreeChanged() override;
 
   std::vector<std::string> GetHttpPaths() const;
@@ -118,6 +119,12 @@ class PrivetHandler : public CloudDelegate::Observer {
   void HandleCheckForUpdates(const base::DictionaryValue& input,
                              const UserInfo& user_info,
                              const RequestCallback& callback);
+  void HandleTraits(const base::DictionaryValue& input,
+                    const UserInfo& user_info,
+                    const RequestCallback& callback);
+  void HandleComponents(const base::DictionaryValue& input,
+                        const UserInfo& user_info,
+                        const RequestCallback& callback);
 
   void ReplyWithSetupStatus(const RequestCallback& callback) const;
   void ReplyToUpdateRequest(const RequestCallback& callback) const;
@@ -139,14 +146,16 @@ class PrivetHandler : public CloudDelegate::Observer {
     RequestCallback callback;
     int request_id = 0;
     int state_fingerprint = -1;
-    int command_defs_fingerprint = -1;
+    int traits_fingerprint = -1;
+    int components_fingerprint = -1;
   };
   std::vector<UpdateRequestParameters> update_requests_;
   int last_update_request_id_{0};
 
   uint64_t last_user_id_{0};
   int state_fingerprint_{0};
-  int command_defs_fingerprint_{0};
+  int traits_fingerprint_{0};
+  int components_fingerprint_{0};
   ScopedObserver<CloudDelegate, CloudDelegate::Observer> cloud_observer_{this};
 
   base::WeakPtrFactory<PrivetHandler> weak_ptr_factory_{this};
diff --git a/src/privet/privet_handler_unittest.cc b/src/privet/privet_handler_unittest.cc
index df4555c..77cabdd 100644
--- a/src/privet/privet_handler_unittest.cc
+++ b/src/privet/privet_handler_unittest.cc
@@ -592,7 +592,7 @@ TEST_F(PrivetHandlerSetupTest, State) {
   EXPECT_JSON_EQ("{'state': {'test': {}}, 'fingerprint': '0'}",
                  HandleRequest("/privet/v3/state", "{}"));
 
-  cloud_.NotifyOnComponentTreeChanged();
+  cloud_.NotifyOnStateChanged();
 
   EXPECT_JSON_EQ("{'state': {'test': {}}, 'fingerprint': '1'}",
                  HandleRequest("/privet/v3/state", "{}"));
@@ -608,6 +608,32 @@ TEST_F(PrivetHandlerSetupTest, CommandsDefs) {
                  HandleRequest("/privet/v3/commandDefs", "{}"));
 }
 
+TEST_F(PrivetHandlerSetupTest, Traits) {
+  EXPECT_JSON_EQ("{'traits': {'test': {}}, 'fingerprint': '0'}",
+                 HandleRequest("/privet/v3/traits", "{}"));
+
+  cloud_.NotifyOnTraitDefsChanged();
+
+  EXPECT_JSON_EQ("{'traits': {'test': {}}, 'fingerprint': '1'}",
+                 HandleRequest("/privet/v3/traits", "{}"));
+}
+
+TEST_F(PrivetHandlerSetupTest, Components) {
+  EXPECT_JSON_EQ("{'components': {'test': {}}, 'fingerprint': '0'}",
+                 HandleRequest("/privet/v3/components", "{}"));
+
+  cloud_.NotifyOnComponentTreeChanged();
+
+  EXPECT_JSON_EQ("{'components': {'test': {}}, 'fingerprint': '1'}",
+                 HandleRequest("/privet/v3/components", "{}"));
+
+  // State change will also change the components fingerprint.
+  cloud_.NotifyOnStateChanged();
+
+  EXPECT_JSON_EQ("{'components': {'test': {}}, 'fingerprint': '2'}",
+                 HandleRequest("/privet/v3/components", "{}"));
+}
+
 TEST_F(PrivetHandlerSetupTest, CommandsExecute) {
   const char kInput[] = "{'name': 'test'}";
   base::DictionaryValue command;
@@ -698,10 +724,13 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_NoInput) {
       .WillOnce(Return(base::TimeDelta::Max()));
   cloud_.NotifyOnTraitDefsChanged();
   cloud_.NotifyOnComponentTreeChanged();
+  cloud_.NotifyOnStateChanged();
   const char kInput[] = "{}";
   const char kExpected[] = R"({
    'commandsFingerprint': '1',
-   'stateFingerprint': '1'
+   'stateFingerprint': '1',
+   'traitsFingerprint': '1',
+   'componentsFingerprint': '2'
   })";
   EXPECT_JSON_EQ(kExpected,
                  HandleRequest("/privet/v3/checkForUpdates", kInput));
@@ -713,13 +742,18 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_AlreadyChanged) {
       .WillOnce(Return(base::TimeDelta::Max()));
   cloud_.NotifyOnTraitDefsChanged();
   cloud_.NotifyOnComponentTreeChanged();
+  cloud_.NotifyOnStateChanged();
   const char kInput[] = R"({
    'commandsFingerprint': '0',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
   })";
   const char kExpected[] = R"({
    'commandsFingerprint': '1',
-   'stateFingerprint': '1'
+   'stateFingerprint': '1',
+   'traitsFingerprint': '1',
+   'componentsFingerprint': '2'
   })";
   EXPECT_JSON_EQ(kExpected,
                  HandleRequest("/privet/v3/checkForUpdates", kInput));
@@ -731,7 +765,9 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_LongPollCommands) {
       .WillOnce(Return(base::TimeDelta::Max()));
   const char kInput[] = R"({
    'commandsFingerprint': '0',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
   })";
   EXPECT_JSON_EQ("{}", HandleRequest("/privet/v3/checkForUpdates", kInput));
   EXPECT_EQ(0, GetResponseCount());
@@ -739,7 +775,31 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_LongPollCommands) {
   EXPECT_EQ(1, GetResponseCount());
   const char kExpected[] = R"({
    'commandsFingerprint': '1',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '1',
+   'componentsFingerprint': '0'
+  })";
+  EXPECT_JSON_EQ(kExpected, GetResponse());
+}
+
+TEST_F(PrivetHandlerSetupTest, CheckForUpdates_LongPollTraits) {
+  EXPECT_CALL(device_, GetHttpRequestTimeout())
+      .WillOnce(Return(base::TimeDelta::Max()));
+  const char kInput[] = R"({
+   'commandsFingerprint': '0',
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
+  })";
+  EXPECT_JSON_EQ("{}", HandleRequest("/privet/v3/checkForUpdates", kInput));
+  EXPECT_EQ(0, GetResponseCount());
+  cloud_.NotifyOnTraitDefsChanged();
+  EXPECT_EQ(1, GetResponseCount());
+  const char kExpected[] = R"({
+   'commandsFingerprint': '1',
+   'stateFingerprint': '0',
+   'traitsFingerprint': '1',
+   'componentsFingerprint': '0'
   })";
   EXPECT_JSON_EQ(kExpected, GetResponse());
 }
@@ -749,7 +809,31 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_LongPollState) {
       .WillOnce(Return(base::TimeDelta::Max()));
   const char kInput[] = R"({
    'commandsFingerprint': '0',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
+  })";
+  EXPECT_JSON_EQ("{}", HandleRequest("/privet/v3/checkForUpdates", kInput));
+  EXPECT_EQ(0, GetResponseCount());
+  cloud_.NotifyOnStateChanged();
+  EXPECT_EQ(1, GetResponseCount());
+  const char kExpected[] = R"({
+   'commandsFingerprint': '0',
+   'stateFingerprint': '1',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '1'
+  })";
+  EXPECT_JSON_EQ(kExpected, GetResponse());
+}
+
+TEST_F(PrivetHandlerSetupTest, CheckForUpdates_LongPollComponents) {
+  EXPECT_CALL(device_, GetHttpRequestTimeout())
+      .WillOnce(Return(base::TimeDelta::Max()));
+  const char kInput[] = R"({
+   'commandsFingerprint': '0',
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
   })";
   EXPECT_JSON_EQ("{}", HandleRequest("/privet/v3/checkForUpdates", kInput));
   EXPECT_EQ(0, GetResponseCount());
@@ -757,16 +841,19 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_LongPollState) {
   EXPECT_EQ(1, GetResponseCount());
   const char kExpected[] = R"({
    'commandsFingerprint': '0',
-   'stateFingerprint': '1'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '1'
   })";
   EXPECT_JSON_EQ(kExpected, GetResponse());
 }
 
-TEST_F(PrivetHandlerSetupTest, CheckForUpdates_LongPollIgnoreCommands) {
+TEST_F(PrivetHandlerSetupTest, CheckForUpdates_LongPollIgnoreTraits) {
   EXPECT_CALL(device_, GetHttpRequestTimeout())
       .WillOnce(Return(base::TimeDelta::Max()));
   const char kInput[] = R"({
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'componentsFingerprint': '0'
   })";
   EXPECT_JSON_EQ("{}", HandleRequest("/privet/v3/checkForUpdates", kInput));
   EXPECT_EQ(0, GetResponseCount());
@@ -776,7 +863,9 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_LongPollIgnoreCommands) {
   EXPECT_EQ(1, GetResponseCount());
   const char kExpected[] = R"({
    'commandsFingerprint': '1',
-   'stateFingerprint': '1'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '1',
+   'componentsFingerprint': '1'
   })";
   EXPECT_JSON_EQ(kExpected, GetResponse());
 }
@@ -785,17 +874,22 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_LongPollIgnoreState) {
   EXPECT_CALL(device_, GetHttpRequestTimeout())
       .WillOnce(Return(base::TimeDelta::Max()));
   const char kInput[] = R"({
-   'commandsFingerprint': '0'
+   'commandsFingerprint': '0',
+   'traitsFingerprint': '0'
   })";
   EXPECT_JSON_EQ("{}", HandleRequest("/privet/v3/checkForUpdates", kInput));
   EXPECT_EQ(0, GetResponseCount());
+  cloud_.NotifyOnStateChanged();
+  EXPECT_EQ(0, GetResponseCount());
   cloud_.NotifyOnComponentTreeChanged();
   EXPECT_EQ(0, GetResponseCount());
   cloud_.NotifyOnTraitDefsChanged();
   EXPECT_EQ(1, GetResponseCount());
   const char kExpected[] = R"({
    'commandsFingerprint': '1',
-   'stateFingerprint': '1'
+   'stateFingerprint': '1',
+   'traitsFingerprint': '1',
+   'componentsFingerprint': '2'
   })";
   EXPECT_JSON_EQ(kExpected, GetResponse());
 }
@@ -806,11 +900,15 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_InstantTimeout) {
   const char kInput[] = R"({
    'commandsFingerprint': '0',
    'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0',
    'waitTimeout': 0
   })";
   const char kExpected[] = R"({
    'commandsFingerprint': '0',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
   })";
   EXPECT_JSON_EQ(kExpected,
                  HandleRequest("/privet/v3/checkForUpdates", kInput));
@@ -822,6 +920,8 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_UserTimeout) {
   const char kInput[] = R"({
    'commandsFingerprint': '0',
    'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0',
    'waitTimeout': 3
   })";
   base::Closure callback;
@@ -833,7 +933,9 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_UserTimeout) {
   EXPECT_EQ(1, GetResponseCount());
   const char kExpected[] = R"({
    'commandsFingerprint': '0',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
   })";
   EXPECT_JSON_EQ(kExpected, GetResponse());
 }
@@ -843,7 +945,9 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_ServerTimeout) {
       .WillOnce(Return(base::TimeDelta::FromMinutes(1)));
   const char kInput[] = R"({
    'commandsFingerprint': '0',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
   })";
   base::Closure callback;
   EXPECT_CALL(device_, PostDelayedTask(_, _, base::TimeDelta::FromSeconds(50)))
@@ -854,7 +958,9 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_ServerTimeout) {
   EXPECT_EQ(1, GetResponseCount());
   const char kExpected[] = R"({
    'commandsFingerprint': '0',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
   })";
   EXPECT_JSON_EQ(kExpected, GetResponse());
 }
@@ -864,7 +970,9 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_VeryShortServerTimeout) {
       .WillOnce(Return(base::TimeDelta::FromSeconds(5)));
   const char kInput[] = R"({
    'commandsFingerprint': '0',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
   })";
   EXPECT_JSON_EQ(kInput, HandleRequest("/privet/v3/checkForUpdates", kInput));
   EXPECT_EQ(1, GetResponseCount());
@@ -876,6 +984,8 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_ServerAndUserTimeout) {
   const char kInput[] = R"({
    'commandsFingerprint': '0',
    'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0',
    'waitTimeout': 10
   })";
   base::Closure callback;
@@ -887,7 +997,9 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_ServerAndUserTimeout) {
   EXPECT_EQ(1, GetResponseCount());
   const char kExpected[] = R"({
    'commandsFingerprint': '0',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0'
   })";
   EXPECT_JSON_EQ(kExpected, GetResponse());
 }
@@ -898,6 +1010,8 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_ChangeBeforeTimeout) {
   const char kInput[] = R"({
    'commandsFingerprint': '0',
    'stateFingerprint': '0',
+   'traitsFingerprint': '0',
+   'componentsFingerprint': '0',
    'waitTimeout': 10
   })";
   base::Closure callback;
@@ -909,7 +1023,9 @@ TEST_F(PrivetHandlerSetupTest, CheckForUpdates_ChangeBeforeTimeout) {
   EXPECT_EQ(1, GetResponseCount());
   const char kExpected[] = R"({
    'commandsFingerprint': '1',
-   'stateFingerprint': '0'
+   'stateFingerprint': '0',
+   'traitsFingerprint': '1',
+   'componentsFingerprint': '0'
   })";
   EXPECT_JSON_EQ(kExpected, GetResponse());
   callback.Run();
diff --git a/src/weave_unittest.cc b/src/weave_unittest.cc
index dd10cbe..b8c4d9b 100644
--- a/src/weave_unittest.cc
+++ b/src/weave_unittest.cc
@@ -45,15 +45,25 @@ using provider::test::MockHttpClientResponse;
 using test::CreateDictionaryValue;
 using test::ValueToString;
 
-const char kCommandDefs[] = R"({
-  "base": {
-    "reboot": {
-      "minimalRole": "user"
+const char kTraitDefs[] = R"({
+  "trait1": {
+    "commands": {
+      "reboot": {
+        "minimalRole": "user"
+      },
+      "shutdown": {
+        "minimalRole": "user",
+        "parameters": {},
+        "results": {}
+      }
     },
-    "_shutdown": {
-      "minimalRole": "user",
-      "parameters": {},
-      "results": {}
+    "state": {
+      "firmwareVersion": {"type": "string"}
+    }
+  },
+  "trait2": {
+    "state": {
+      "battery_level": {"type": "integer"}
     }
   }
 })";
@@ -72,7 +82,7 @@ const char kDeviceResource[] = R"({
   "description": "Developer device",
   "stateValidationEnabled": true,
   "commandDefs":{
-    "base": {
+    "trait1": {
       "reboot": {
         "minimalRole": "user",
         "parameters": {"delay": {"type": "integer"}},
@@ -86,15 +96,39 @@ const char kDeviceResource[] = R"({
     }
   },
   "state":{
-    "base":{
-      "firmwareVersion":"FIRMWARE_VERSION",
-      "localAnonymousAccessMaxRole":"viewer",
-      "localDiscoveryEnabled":true,
-      "localPairingEnabled":true,
-      "network":{
+    "trait1": {"firmwareVersion":"FIRMWARE_VERSION"},
+    "trait2": {"battery_level":44}
+  },
+  "traits": {
+    "trait1": {
+      "commands": {
+        "reboot": {
+          "minimalRole": "user"
+        },
+        "shutdown": {
+          "minimalRole": "user",
+          "parameters": {},
+          "results": {}
+        }
+      },
+      "state": {
+        "firmwareVersion": {"type": "string"}
       }
     },
-    "power": {"battery_level":44}
+    "trait2": {
+      "state": {
+        "battery_level": {"type": "integer"}
+      }
+    }
+  },
+  "components": {
+    "myComponent": {
+      "traits": ["trait1", "trait2"],
+      "state": {
+        "trait1": {"firmwareVersion":"FIRMWARE_VERSION"},
+        "trait2": {"battery_level":44}
+      }
+    }
   }
 })";
 
@@ -127,10 +161,6 @@ const char kAuthTokenResponse[] = R"({
   "refresh_token" : "REFRESH_TOKEN"
 })";
 
-const char kStateDefs[] = R"({"power": {"battery_level":"integer"}})";
-
-const char kStateDefaults[] = R"({"power": {"battery_level":44}})";
-
 MATCHER_P(MatchTxt, txt, "") {
   std::vector<std::string> txt_copy = txt;
   std::sort(txt_copy.begin(), txt_copy.end());
@@ -259,14 +289,17 @@ class WeaveTest : public ::testing::Test {
                   "/privet/v3/checkForUpdates", "/privet/v3/commandDefs",
                   "/privet/v3/commands/cancel", "/privet/v3/commands/execute",
                   "/privet/v3/commands/list", "/privet/v3/commands/status",
-                  "/privet/v3/pairing/cancel", "/privet/v3/pairing/confirm",
-                  "/privet/v3/pairing/start", "/privet/v3/setup/start",
-                  "/privet/v3/setup/status", "/privet/v3/state"}),
+                  "/privet/v3/components", "/privet/v3/pairing/cancel",
+                  "/privet/v3/pairing/confirm", "/privet/v3/pairing/start",
+                  "/privet/v3/setup/start", "/privet/v3/setup/status",
+                  "/privet/v3/state", "/privet/v3/traits"}),
               GetKeys(https_handlers_));
 
-    device_->AddCommandDefinitionsFromJson(kCommandDefs);
-    device_->AddStateDefinitionsFromJson(kStateDefs);
-    device_->SetStatePropertiesFromJson(kStateDefaults, nullptr);
+    device_->AddTraitDefinitionsFromJson(kTraitDefs);
+    EXPECT_TRUE(device_->AddComponent("myComponent", {"trait1", "trait2"},
+                                      nullptr));
+    EXPECT_TRUE(device_->SetStatePropertiesFromJson(
+        "myComponent", R"({"trait2": {"battery_level":44}})", nullptr));
 
     task_runner_.Run();
   }
@@ -324,7 +357,9 @@ TEST_F(WeaveTest, StartNoWifi) {
   device_ = weave::Device::Create(&config_store_, &task_runner_, &http_client_,
                                   &network_, &dns_sd_, &http_server_, nullptr,
                                   &bluetooth_);
-  device_->AddCommandDefinitionsFromJson(kCommandDefs);
+  device_->AddTraitDefinitionsFromJson(kTraitDefs);
+  EXPECT_TRUE(device_->AddComponent("myComponent", {"trait1", "trait2"},
+                                    nullptr));
 
   task_runner_.Run();
 }