Allow re-enabling feature Prefetch Profiles

Summary: For the prefetch-profile feature, added a matching enable option to go with the existing disable option. Enable temporarily enables the prefetch-profile feature for cases where a user has already used disable.

Reviewed By: kmancini

Differential Revision: D27191711

fbshipit-source-id: cf9cb7d4054e9fcee66bdb5f205562ec79c2757e

eden/fs/cli/prefetch_profile.py

@@ -525,17 +525,8 @@ class FetchProfileCmd(Subcmd):
     "Disables prefetch profiles locally",
 )
 class DisableProfileCmd(Subcmd):
-    def setup_parser(self, parser: argparse.ArgumentParser) -> None:
-        parser.add_argument(
-            "--checkout",
-            help="The checkout for which prefetching should be disabled",
-            default=None,
-        )
-
     def run(self, args: argparse.Namespace) -> int:
-        checkout = args.checkout
+        instance = get_eden_instance(args)
-
-        instance, _checkout, _rel_path = require_checkout(args, checkout)
         config = instance.read_local_config()
         prefetch_profiles_section = {}
         if config.has_section("prefetch-profiles"):
@@ -549,6 +540,26 @@ class DisableProfileCmd(Subcmd):
         return 0
 
 
+@prefetch_profile_cmd(
+    "enable",
+    "Enables prefetch profiles locally",
+)
+class EnableProfileCmd(Subcmd):
+    def run(self, args: argparse.Namespace) -> int:
+        instance = get_eden_instance(args)
+        config = instance.read_local_config()
+        prefetch_profiles_section = {}
+        if config.has_section("prefetch-profiles"):
+            prefetch_profiles_section.update(
+                config.get_section_str_to_any("prefetch-profiles")
+            )
+        prefetch_profiles_section["prefetching-enabled"] = True
+        config["prefetch-profiles"] = prefetch_profiles_section
+        instance.write_local_config(config)
+
+        return 0
+
+
 class PrefetchProfileCmd(Subcmd):
     NAME = "prefetch-profile"
     HELP = (

eden/fs/cli/test/prefetch_profile_test.py

@@ -13,6 +13,7 @@ import unittest.mock as mock
 from eden.fs.cli.config import EdenInstance
 from eden.fs.cli.configutil import EdenConfigParser
 from eden.fs.cli.prefetch_profile import DisableProfileCmd
+from eden.fs.cli.prefetch_profile import EnableProfileCmd
 
 
 class PrefetchProfileTest(unittest.TestCase):
@@ -21,10 +22,9 @@ class PrefetchProfileTest(unittest.TestCase):
         self.mock_argument_parser = mock.MagicMock(spec=argparse.ArgumentParser)
         self.mock_args = mock.MagicMock(spec=argparse.Namespace)
 
-    @mock.patch("eden.fs.cli.prefetch_profile.require_checkout")
+    @mock.patch("eden.fs.cli.prefetch_profile.get_eden_instance")
-    def test_disable_no_config(self, mock_require_checkout: mock.MagicMock) -> None:
+    def test_disable_no_config(self, mock_get_eden_instance: mock.MagicMock) -> None:
-        self.mock_args.checkout = "test"
+        mock_get_eden_instance.return_value = self.mock_eden_instance
-        mock_require_checkout.return_value = (self.mock_eden_instance, None, None)
 
         test_config = EdenConfigParser()
         self.mock_eden_instance.read_local_config.return_value = test_config
@@ -39,12 +39,11 @@ class PrefetchProfileTest(unittest.TestCase):
             {"prefetching-enabled": False},
         )
 
-    @mock.patch("eden.fs.cli.prefetch_profile.require_checkout")
+    @mock.patch("eden.fs.cli.prefetch_profile.get_eden_instance")
     def test_disable_existing_config(
-        self, mock_require_checkout: mock.MagicMock
+        self, mock_get_eden_instance: mock.MagicMock
     ) -> None:
-        self.mock_args.checkout = "test"
+        mock_get_eden_instance.return_value = self.mock_eden_instance
-        mock_require_checkout.return_value = (self.mock_eden_instance, None, None)
 
         test_config = EdenConfigParser()
         test_config.read_dict(
@@ -72,3 +71,53 @@ class PrefetchProfileTest(unittest.TestCase):
         self.assertEqual(
             test_config.get_section_str_to_any("something-random"), {"yup": "test"}
         )
+
+    @mock.patch("eden.fs.cli.prefetch_profile.get_eden_instance")
+    def test_enable_no_config(self, mock_get_eden_instance: mock.MagicMock) -> None:
+        mock_get_eden_instance.return_value = self.mock_eden_instance
+
+        test_config = EdenConfigParser()
+        self.mock_eden_instance.read_local_config.return_value = test_config
+
+        test_enable = EnableProfileCmd(self.mock_argument_parser)
+        test_enable.run(self.mock_args)
+
+        self.mock_eden_instance.read_local_config.assert_called_once()
+        self.mock_eden_instance.write_local_config.assert_called_once_with(test_config)
+        self.assertEqual(
+            test_config.get_section_str_to_any("prefetch-profiles"),
+            {"prefetching-enabled": True},
+        )
+
+    @mock.patch("eden.fs.cli.prefetch_profile.get_eden_instance")
+    def test_enable_existing_config(
+        self, mock_get_eden_instance: mock.MagicMock
+    ) -> None:
+        mock_get_eden_instance.return_value = self.mock_eden_instance
+
+        test_config = EdenConfigParser()
+        test_config.read_dict(
+            {
+                "prefetch-profiles": {
+                    "prefetching-enabled": False,
+                },
+                "something-random": {
+                    "yup": "test",
+                },
+            }
+        )
+        self.mock_eden_instance.read_local_config.return_value = test_config
+
+        test_enable = EnableProfileCmd(self.mock_argument_parser)
+        test_enable.run(self.mock_args)
+
+        self.mock_eden_instance.read_local_config.assert_called_once()
+        self.mock_eden_instance.write_local_config.assert_called_once_with(test_config)
+        self.assertEqual(
+            test_config.get_section_str_to_any("prefetch-profiles"),
+            {"prefetching-enabled": True},
+        )
+        self.assertTrue(test_config.has_section("something-random"))
+        self.assertEqual(
+            test_config.get_section_str_to_any("something-random"), {"yup": "test"}
+        )