moved dbgym_cfg to the util file

cmu-db · Dec 21, 2024 · 4f7362b · 4f7362b
1 parent f19c3ae
commit 4f7362b
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 17 deletions.
diff --git a/env/integtest_pg_conn.py b/env/integtest_pg_conn.py
@@ -1,12 +1,10 @@
 import copy
 import subprocess
 import unittest
-from pathlib import Path
-
-import yaml
 
 from env.integtest_util import (
     ENV_INTEGTESTS_DBGYM_CONFIG_FPATH,
+    INTEGTEST_DBGYM_CFG,
     get_integtest_workspace_path,
 )
 from env.pg_conn import PostgresConn
@@ -36,28 +34,26 @@ def setUpClass() -> None:
         if not get_integtest_workspace_path().exists():
             subprocess.run(["./env/set_up_env_integtests.sh"], check=True)
 
-        PostgresConnTests.dbgym_cfg = DBGymConfig(ENV_INTEGTESTS_DBGYM_CONFIG_FPATH)
-
     def setUp(self) -> None:
         self.assertFalse(
             get_is_postgres_running(),
             "Make sure Postgres isn't running before starting the integration test. `pkill postgres` is one way "
             + "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(
-            self.dbgym_cfg.dbgym_workspace_path, BENCHMARK, SCALE_FACTOR
+            INTEGTEST_DBGYM_CFG.dbgym_workspace_path, BENCHMARK, SCALE_FACTOR
         )
         self.dbdata_parent_dpath = default_dbdata_parent_dpath(
-            self.dbgym_cfg.dbgym_workspace_path
+            INTEGTEST_DBGYM_CFG.dbgym_workspace_path
         )
-        self.pgbin_dpath = default_pgbin_path(self.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(
-            PostgresConnTests.dbgym_cfg,
+            INTEGTEST_DBGYM_CFG,
             pgport,
             self.pristine_dbdata_snapshot_path,
             self.dbdata_parent_dpath,

diff --git a/env/integtest_tuning_agent.py b/env/integtest_tuning_agent.py
@@ -3,11 +3,10 @@
 from typing import Any, Optional
 
 from env.integtest_util import (
-    ENV_INTEGTESTS_DBGYM_CONFIG_FPATH,
+    INTEGTEST_DBGYM_CFG,
     get_integtest_workspace_path,
 )
 from env.tuning_agent import DBMSConfigDelta, TuningAgent
-from util.workspace import DBGymConfig
 
 
 class MockTuningAgent(TuningAgent):
@@ -24,22 +23,18 @@ def _step(self) -> DBMSConfigDelta:
 
 
 class PostgresConnTests(unittest.TestCase):
-    dbgym_cfg: DBGymConfig
-
     @staticmethod
     def setUpClass() -> None:
         # If you're running the test locally, this check makes runs past the first one much faster.
         if not get_integtest_workspace_path().exists():
             subprocess.run(["./env/set_up_env_integtests.sh"], check=True)
 
-        PostgresConnTests.dbgym_cfg = DBGymConfig(ENV_INTEGTESTS_DBGYM_CONFIG_FPATH)
-
     @staticmethod
     def make_config(letter: str) -> DBMSConfigDelta:
         return DBMSConfigDelta([letter], {letter: letter}, {letter: [letter]})
 
     def test_get_step_delta(self) -> None:
-        agent = MockTuningAgent(PostgresConnTests.dbgym_cfg)
+        agent = MockTuningAgent(INTEGTEST_DBGYM_CFG)
 
         agent.config_to_return = PostgresConnTests.make_config("a")
         agent.step()
@@ -54,7 +49,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(PostgresConnTests.dbgym_cfg)
+        agent = MockTuningAgent(INTEGTEST_DBGYM_CFG)
 
         agent.config_to_return = PostgresConnTests.make_config("a")
         agent.step()

diff --git a/env/integtest_util.py b/env/integtest_util.py
@@ -2,7 +2,11 @@
 
 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)
 
 
 def get_integtest_workspace_path() -> Path: