From 4f7362b8b128ad18227794daff7a9d4b9e13fc92 Mon Sep 17 00:00:00 2001
From: Patrick Wang <wang.patrick57@gmail.com>
Date: Fri, 20 Dec 2024 20:34:50 -0500
Subject: [PATCH] moved dbgym_cfg to the util file

---
 env/integtest_pg_conn.py      | 14 +++++---------
 env/integtest_tuning_agent.py | 11 +++--------
 env/integtest_util.py         |  4 ++++
 3 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/env/integtest_pg_conn.py b/env/integtest_pg_conn.py
index 39364119..9eff9040 100644
--- 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,8 +34,6 @@ 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(),
@@ -45,19 +41,19 @@ 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(
-            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
index 35ca329a..081e8281 100644
--- 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
index 4cf185e8..d48b4f58 100644
--- 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: