integtest workspace now creates dbgym config after creating workspace

env/integtest_pg_conn.py

@@ -1,13 +1,7 @@
 import copy
-import subprocess
 import unittest
 
-from env.integtest_util import (
-    ENV_INTEGTESTS_DBGYM_CONFIG_FPATH,
-    INTEGTEST_DBGYM_CFG,
-    get_integtest_workspace_path,
-    set_up_integtest_workspace,
-)
+from env.integtest_util import IntegtestWorkspace
 from env.pg_conn import PostgresConn
 from util.pg import (
     DEFAULT_POSTGRES_PORT,
@@ -31,7 +25,7 @@ class PostgresConnTests(unittest.TestCase):
 
     @staticmethod
     def setUpClass() -> None:
-        set_up_integtest_workspace()
+        IntegtestWorkspace.set_up_workspace()
 
     def setUp(self) -> None:
         self.assertFalse(
@@ -40,19 +34,23 @@ def setUp(self) -> None:
             + "to ensure this. Be careful about accidentally taking down other people's Postgres instances though.",
         )
         self.pristine_dbdata_snapshot_path = default_pristine_dbdata_snapshot_path(
-            INTEGTEST_DBGYM_CFG.dbgym_workspace_path, BENCHMARK, SCALE_FACTOR
+            IntegtestWorkspace.get_dbgym_cfg().dbgym_workspace_path,
+            BENCHMARK,
+            SCALE_FACTOR,
         )
         self.dbdata_parent_dpath = default_dbdata_parent_dpath(
-            INTEGTEST_DBGYM_CFG.dbgym_workspace_path
+            IntegtestWorkspace.get_dbgym_cfg().dbgym_workspace_path
+        )
+        self.pgbin_dpath = default_pgbin_path(
+            IntegtestWorkspace.get_dbgym_cfg().dbgym_workspace_path
         )
-        self.pgbin_dpath = default_pgbin_path(INTEGTEST_DBGYM_CFG.dbgym_workspace_path)
 
     def tearDown(self) -> None:
         self.assertFalse(get_is_postgres_running())
 
     def create_pg_conn(self, pgport: int = DEFAULT_POSTGRES_PORT) -> PostgresConn:
         return PostgresConn(
-            INTEGTEST_DBGYM_CFG,
+            IntegtestWorkspace.get_dbgym_cfg(),
             pgport,
             self.pristine_dbdata_snapshot_path,
             self.dbdata_parent_dpath,

env/tmprenamed_tuning_agent.py → env/integtest_tuning_agent.py

@@ -1,12 +1,7 @@
-import subprocess
 import unittest
 from typing import Any, Optional
 
-from env.integtest_util import (
-    INTEGTEST_DBGYM_CFG,
-    get_integtest_workspace_path,
-    set_up_integtest_workspace,
-)
+from env.integtest_util import IntegtestWorkspace
 from env.tuning_agent import DBMSConfigDelta, TuningAgent
 
 
@@ -26,14 +21,14 @@ def _step(self) -> DBMSConfigDelta:
 class PostgresConnTests(unittest.TestCase):
     @staticmethod
     def setUpClass() -> None:
-        set_up_integtest_workspace()
+        IntegtestWorkspace.set_up_workspace()
 
     @staticmethod
     def make_config(letter: str) -> DBMSConfigDelta:
         return DBMSConfigDelta([letter], {letter: letter}, {letter: [letter]})
 
     def test_get_step_delta(self) -> None:
-        agent = MockTuningAgent(INTEGTEST_DBGYM_CFG)
+        agent = MockTuningAgent(IntegtestWorkspace.get_dbgym_cfg())
 
         agent.config_to_return = PostgresConnTests.make_config("a")
         agent.step()
@@ -48,7 +43,7 @@ def test_get_step_delta(self) -> None:
         self.assertEqual(agent.get_step_delta(2), PostgresConnTests.make_config("c"))
 
     def test_get_all_deltas(self) -> None:
-        agent = MockTuningAgent(INTEGTEST_DBGYM_CFG)
+        agent = MockTuningAgent(IntegtestWorkspace.get_dbgym_cfg())
 
         agent.config_to_return = PostgresConnTests.make_config("a")
         agent.step()

env/integtest_util.py

@@ -1,21 +1,40 @@
 import subprocess
 from pathlib import Path
+from typing import Optional
 
 import yaml
 
 from util.workspace import DBGymConfig
 
-ENV_INTEGTESTS_DBGYM_CONFIG_FPATH = Path("env/env_integtests_dbgym_config.yaml")
-INTEGTEST_DBGYM_CFG = DBGymConfig(ENV_INTEGTESTS_DBGYM_CONFIG_FPATH)
 
+class IntegtestWorkspace:
+    """
+    This is essentially a singleton class. This avoids multiple integtest_*.py files creating
+    the workspace and/or the DBGymConfig redundantly.
+    """
 
-def set_up_integtest_workspace() -> None:
-    # This if statement prevents us from setting up the workspace twice, which saves time.
-    if not get_integtest_workspace_path().exists():
-        subprocess.run(["./env/set_up_env_integtests.sh"], check=True)
+    ENV_INTEGTESTS_DBGYM_CONFIG_FPATH = Path("env/env_integtests_dbgym_config.yaml")
+    INTEGTEST_DBGYM_CFG: Optional[DBGymConfig] = None
 
+    @staticmethod
+    def set_up_workspace() -> None:
+        # This if statement prevents us from setting up the workspace twice, which saves time.
+        if not IntegtestWorkspace.get_workspace_path().exists():
+            subprocess.run(["./env/set_up_env_integtests.sh"], check=True)
 
-def get_integtest_workspace_path() -> Path:
-    with open(ENV_INTEGTESTS_DBGYM_CONFIG_FPATH) as f:
-        return Path(yaml.safe_load(f)["dbgym_workspace_path"])
-    assert False
+        # The DBGymConfig needs to be created after running ./env/set_up_env_integtests.sh so
+        # that it is created correctly.
+        IntegtestWorkspace.INTEGTEST_DBGYM_CFG = DBGymConfig(
+            IntegtestWorkspace.ENV_INTEGTESTS_DBGYM_CONFIG_FPATH
+        )
+
+    @staticmethod
+    def get_dbgym_cfg() -> DBGymConfig:
+        assert IntegtestWorkspace.INTEGTEST_DBGYM_CFG is not None
+        return IntegtestWorkspace.INTEGTEST_DBGYM_CFG
+
+    @staticmethod
+    def get_workspace_path() -> Path:
+        with open(IntegtestWorkspace.ENV_INTEGTESTS_DBGYM_CONFIG_FPATH) as f:
+            return Path(yaml.safe_load(f)["dbgym_workspace_path"])
+        assert False

util/workspace.py

@@ -10,7 +10,7 @@
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
-from typing import IO, Any, Callable, Optional, Tuple
+from typing import IO, Any, Callable, Optional
 
 import redis
 import yaml