From bf4f130a48c5c840ab4423ddc5c60055b0fa5f8b Mon Sep 17 00:00:00 2001
From: dachengx <dx2227@columbia.edu>
Date: Tue, 16 Jan 2024 02:05:38 -0600
Subject: [PATCH] Bug fixes

---
 tests/test_contexts.py | 2 +-
 tests/test_wfsim.py    | 2 +-
 wfsim/__init__.py      | 2 ++
 wfsim/contexts.py      | 8 ++++++--
 4 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/tests/test_contexts.py b/tests/test_contexts.py
index 86666fe6..43b30945 100644
--- a/tests/test_contexts.py
+++ b/tests/test_contexts.py
@@ -78,7 +78,7 @@ def test_nt_sim_context_bad_inits(self):
 
 
 def test_sim_context():
-    straxen.contexts.xenon1t_simulation()
+    wfsim.contexts.xenon1t_simulation()
 
 
 @skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!")
diff --git a/tests/test_wfsim.py b/tests/test_wfsim.py
index 4e2c8dd4..623ac5e5 100644
--- a/tests/test_wfsim.py
+++ b/tests/test_wfsim.py
@@ -51,7 +51,7 @@ def test_sim_1T():
                     **conf_1t,
                 ),
                 **straxen.legacy.x1t_common_config),
-            **straxen.legacy.contexts_1t.get_x1t_context_config(),
+            **straxen.legacy.get_x1t_context_config(),
         )
         st.register(wfsim.RawRecordsFromFax1T)
         log.debug(f'Setting testing config {testing_config_1t}')
diff --git a/wfsim/__init__.py b/wfsim/__init__.py
index c2197ebb..9c3fa797 100644
--- a/wfsim/__init__.py
+++ b/wfsim/__init__.py
@@ -11,3 +11,5 @@
 
 from .load_resource import *
 from .utils import *
+
+from .contexts import *
diff --git a/wfsim/contexts.py b/wfsim/contexts.py
index e062de3a..c3a48586 100644
--- a/wfsim/contexts.py
+++ b/wfsim/contexts.py
@@ -2,6 +2,7 @@
 from immutabledict import immutabledict
 import strax
 import straxen
+import wfsim
 from .strax_interface import RawRecordsFromFax1T
 
 
@@ -280,8 +281,11 @@ def xenonnt_simulation(
 def xenon1t_simulation(output_folder="./strax_data"):
     st = strax.Context(
         storage=strax.DataDirectory(output_folder),
-        config=dict(fax_config="fax_config_1t.json", detector="XENON1T", **straxen.contexts.x1t_common_config),
-        **straxen.contexts.get_x1t_context_config(),
+        config=dict(
+            fax_config="fax_config_1t.json", detector="XENON1T",
+            **straxen.legacy.x1t_common_config,
+        ),
+        **straxen.legacy.get_x1t_context_config(),
     )
     st.register(RawRecordsFromFax1T)
     st.deregister_plugins_with_missing_dependencies()