Fixed vpn button still has connected state for not purchased user

fix brave/brave-browser#33023

When purchased state is changed to *not* purchased during the runtime,
button should not show as connected.

browser/ui/BUILD.gn

@@ -1096,8 +1096,10 @@ source_set("unit_tests") {
       sources += [ "views/toolbar/brave_vpn_button_unittest.cc" ]
       deps += [
         "//base",
+        "//brave/components/brave_vpn/browser",
         "//brave/components/brave_vpn/browser/connection:api",
         "//brave/components/brave_vpn/browser/connection/ikev2:sim",
+        "//brave/components/skus/browser",
         "//chrome/test:test_support",
         "//skia",
         "//testing/gtest",

browser/ui/views/toolbar/brave_vpn_button.cc

@@ -146,6 +146,7 @@ BraveVPNButton::BraveVPNButton(Browser* browser)
       service_(brave_vpn::BraveVpnServiceFactory::GetForProfile(
           browser_->profile())) {
   CHECK(service_);
+  UpdateButtonState();
   Observe(service_);
 
   // Replace ToolbarButton's highlight path generator.
@@ -211,15 +212,14 @@ void BraveVPNButton::OnConnectionStateChanged(ConnectionState state) {
 
 void BraveVPNButton::UpdateButtonState() {
   is_error_state_ = IsConnectError();
+  is_connected_ = IsConnected();
 }
 
 void BraveVPNButton::OnPurchasedStateChanged(
     brave_vpn::mojom::PurchasedState state,
     const absl::optional<std::string>& description) {
   UpdateButtonState();
-  if (IsPurchased()) {
-    UpdateColorsAndInsets();
-  }
+  UpdateColorsAndInsets();
 }
 
 std::unique_ptr<views::Border> BraveVPNButton::GetBorder(
@@ -256,13 +256,12 @@ void BraveVPNButton::UpdateColorsAndInsets() {
     image()->SetBackground(std::make_unique<ConnectErrorIconBackground>(
         cp->GetColor(kColorBraveVpnButtonIconErrorInner)));
   } else {
-    const bool is_connected = IsConnected();
-    SetImage(
-        views::Button::STATE_NORMAL,
-        gfx::CreateVectorIcon(
-            is_connected ? kVpnIndicatorOnIcon : kVpnIndicatorOffIcon,
-            cp->GetColor(is_connected ? kColorBraveVpnButtonIconConnected
-                                      : kColorBraveVpnButtonIconDisconnected)));
+    SetImage(views::Button::STATE_NORMAL,
+             gfx::CreateVectorIcon(
+                 is_connected_ ? kVpnIndicatorOnIcon : kVpnIndicatorOffIcon,
+                 cp->GetColor(is_connected_
+                                  ? kColorBraveVpnButtonIconConnected
+                                  : kColorBraveVpnButtonIconDisconnected)));
 
     // Use background for inner color of button image.
     // Adjusted border thickness to make invisible to the outside of the icon.

browser/ui/views/toolbar/brave_vpn_button.h

@@ -66,6 +66,7 @@ class BraveVPNButton : public ToolbarButton,
   void UpdateButtonState();
 
   bool is_error_state_ = false;
+  bool is_connected_ = false;
   absl::optional<brave_vpn::mojom::ConnectionState>
       connection_state_for_testing_;
   raw_ptr<Browser> browser_ = nullptr;

browser/ui/views/toolbar/brave_vpn_button_unittest.cc

@@ -11,8 +11,10 @@
 #include "base/run_loop.h"
 #include "brave/browser/brave_browser_process_impl.h"
 #include "brave/browser/brave_vpn/brave_vpn_service_factory.h"
+#include "brave/components/brave_vpn/browser/brave_vpn_service.h"
 #include "brave/components/brave_vpn/browser/connection/brave_vpn_os_connection_api.h"
 #include "brave/components/brave_vpn/browser/connection/ikev2/brave_vpn_ras_connection_api_sim.h"
+#include "brave/components/skus/browser/skus_utils.h"
 #include "brave/test/base/testing_brave_browser_process.h"
 #include "chrome/browser/prefs/browser_prefs.h"
 #include "chrome/browser/themes/theme_service_factory.h"
@@ -78,13 +80,30 @@ class BraveVpnButtonUnitTest : public testing::Test {
     profile_.reset();
     base::RunLoop().RunUntilIdle();
   }
+
   bool IsErrorState() { return button_->IsErrorState(); }
-  void FireVpnState(brave_vpn::mojom::ConnectionState state) {
+
+  void FireVpnState(mojom::ConnectionState state) {
     button_->SetVpnConnectionStateForTesting(state);
     button_->OnConnectionStateChanged(state);
   }
 
- private:
+  // Give button's visual connection state.
+  bool DoesButtonHaveConnectedState() const { return button_->is_connected_; }
+
+  void SetPurchasedState(const std::string& env, mojom::PurchasedState state) {
+    button_->service_->SetPurchasedState(env, state);
+  }
+
+  void SetConnectionState(mojom::ConnectionState state) {
+    button_->service_->connection_api_->SetConnectionStateForTesting(state);
+  }
+
+  bool IsOsVpnConnected() const {
+    return button_->service_->GetConnectionState() ==
+           mojom::ConnectionState::CONNECTED;
+  }
+
   content::BrowserTaskEnvironment task_environment_;
   std::unique_ptr<BraveVPNButton> button_;
   ChromeLayoutProvider layout_provider_;
@@ -113,4 +132,29 @@ TEST_F(BraveVpnButtonUnitTest, SkipAttemptsToConnectInFailedState) {
   EXPECT_FALSE(IsErrorState());
 }
 
+TEST_F(BraveVpnButtonUnitTest, ButtonStateTestWithPurchasedState) {
+  // Set underlying vpn state as connected state and check button's connect
+  // state.
+  std::string env = skus::GetDefaultEnvironment();
+  SetConnectionState(mojom::ConnectionState::CONNECTED);
+  EXPECT_TRUE(IsOsVpnConnected());
+
+  // Button should have not connected state when not purchased.
+  SetPurchasedState(env, mojom::PurchasedState::NOT_PURCHASED);
+  base::RunLoop().RunUntilIdle();
+  EXPECT_FALSE(DoesButtonHaveConnectedState());
+  EXPECT_TRUE(IsOsVpnConnected());
+
+  // Button should have connected state when changed to purchased.
+  SetPurchasedState(env, mojom::PurchasedState::PURCHASED);
+  base::RunLoop().RunUntilIdle();
+  EXPECT_TRUE(DoesButtonHaveConnectedState());
+  EXPECT_TRUE(IsOsVpnConnected());
+
+  SetPurchasedState(env, mojom::PurchasedState::NOT_PURCHASED);
+  base::RunLoop().RunUntilIdle();
+  EXPECT_FALSE(DoesButtonHaveConnectedState());
+  EXPECT_TRUE(IsOsVpnConnected());
+}
+
 }  // namespace brave_vpn

browser/ui/webui/skus_internals_ui.cc

@@ -31,6 +31,8 @@
 #include "ui/base/clipboard/scoped_clipboard_writer.h"
 
 #if BUILDFLAG(ENABLE_BRAVE_VPN)
+#include "brave/browser/brave_vpn/brave_vpn_service_factory.h"
+#include "brave/components/brave_vpn/browser/brave_vpn_service.h"
 #include "brave/components/brave_vpn/browser/brave_vpn_service_helper.h"
 #include "brave/components/brave_vpn/browser/connection/brave_vpn_os_connection_api.h"
 #include "brave/components/brave_vpn/common/brave_vpn_utils.h"
@@ -156,6 +158,10 @@ void SkusInternalsUI::ResetSkusState() {
   auto* profile = Profile::FromWebUI(web_ui());
   if (brave_vpn::IsBraveVPNEnabled(profile->GetPrefs())) {
     brave_vpn::ClearSubscriberCredential(local_state_);
+    if (auto* service =
+            brave_vpn::BraveVpnServiceFactory::GetForProfile(profile)) {
+      service->ReloadPurchasedState();
+    }
   }
 #endif
 

components/brave_vpn/browser/brave_vpn_service.h

@@ -157,6 +157,7 @@ class BraveVpnService :
 
  private:
   friend class BraveVPNServiceTest;
+  friend class BraveVpnButtonUnitTest;
 
 #if !BUILDFLAG(IS_ANDROID)
   friend class ::BraveAppMenuBrowserTest;

components/brave_vpn/browser/connection/brave_vpn_os_connection_api.h

@@ -98,6 +98,8 @@ class BraveVPNOSConnectionAPI
       net::NetworkChangeNotifier::ConnectionType type) override;
 
  private:
+  friend class BraveVpnButtonUnitTest;
+
   FRIEND_TEST_ALL_PREFIXES(BraveVPNOSConnectionAPIUnitTest, NeedsConnectTest);
   FRIEND_TEST_ALL_PREFIXES(BraveVPNOSConnectionAPIUnitTest, ConnectionInfoTest);
   FRIEND_TEST_ALL_PREFIXES(BraveVPNOSConnectionAPIUnitTest,