Allow set_priority(0) as alternative for unset_priority

rcaelers · Sep 17, 2024 · 7461c4f · 7461c4f
1 parent 8fc82bc
commit 7461c4f
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 19 deletions.
diff --git a/src/UpgradeControl.cc b/src/UpgradeControl.cc
@@ -139,23 +139,6 @@ UpgradeControl::set_certificate(const std::string &cert)
   http->options().add_ca_cert(cert);
 }
 
-outcome::std_result<void>
-UpgradeControl::set_priority(int priority)
-{
-  if (priority < 1 || priority > 100)
-    {
-      return outcome::failure(unfold::UnfoldErrc::InvalidArgument);
-    }
-  custom_priority = priority;
-  return outcome::success();
-}
-
-void
-UpgradeControl::unset_priority()
-{
-  custom_priority.reset();
-}
-
 void
 UpgradeControl::set_periodic_update_check_enabled(bool enabled)
 {
@@ -389,6 +372,32 @@ UpgradeControl::init_priority()
     }
 }
 
+outcome::std_result<void>
+UpgradeControl::set_priority(int priority)
+{
+  if (priority < 0 || priority > 100)
+    {
+      return outcome::failure(unfold::UnfoldErrc::InvalidArgument);
+    }
+
+  if (priority == 0)
+    {
+      unset_priority();
+    }
+  else
+    {
+      custom_priority = priority;
+    }
+
+  return outcome::success();
+}
+
+void
+UpgradeControl::unset_priority()
+{
+  custom_priority.reset();
+}
+
 int
 UpgradeControl::get_priority() const
 {

diff --git a/test/UpgradeControlTests.cc b/test/UpgradeControlTests.cc
@@ -1029,10 +1029,16 @@ TEST_F(UpgradeControlTest, Priority)
   EXPECT_EQ(control->get_priority(), 5);
   control->unset_priority();
   EXPECT_EQ(control->get_priority(), 10);
+
+  ret = control->set_priority(80);
+  EXPECT_EQ(ret.has_error(), false);
   ret = control->set_priority(101);
   EXPECT_EQ(ret.has_error(), true);
-  ret = control->set_priority(0);
-  EXPECT_EQ(ret.has_error(), true);
   ret = control->set_priority(-1);
   EXPECT_EQ(ret.has_error(), true);
+  EXPECT_EQ(control->get_priority(), 80);
+
+  ret = control->set_priority(0);
+  EXPECT_EQ(ret.has_error(), false);
+  EXPECT_EQ(control->get_priority(), 10);
 }