fix test

Signed-off-by: kerthcet <kerthcet@gmail.com>

Author: kerthcet

.env.test

@@ -3,4 +3,5 @@ ALPHATRION_METADATA_INIT_TABLES=true
 ALPHATRION_ARTIFACT_REGISTRY_URL=localhost:5001
 ALPHATRION_ARTIFACT_INSECURE=true
 ALPHATRION_LOG_LEVEL=INFO
-ALPHATRION_AUTO_CLEANUP=true
+ALPHATRION_AUTO_CLEANUP=true
+ALPHATRION_ENABLE_TRACING=true

alphatrion/storage/runtime.py

@@ -22,6 +22,10 @@ def __init__(self, force_init: bool = False):
         if self._inited and not force_init:
             return
 
+        # Reset inited flag when force_init is True
+        if force_init:
+            self._inited = False
+
         self._metadb = SQLStore(
             os.getenv(envs.METADATA_DB_URL),
             init_tables=os.getenv(envs.METADATA_INIT_TABLES, "false").lower() == "true",
@@ -67,7 +71,7 @@ def init(force_init: bool = False):
     """
 
     global __STORAGE_RUNTIME__
-    if __STORAGE_RUNTIME__ is None:
+    if force_init or __STORAGE_RUNTIME__ is None:
         __STORAGE_RUNTIME__ = StorageRuntime(force_init=force_init)
 
 

tests/integration/server/test_graphql_query.py

@@ -396,6 +396,10 @@ async def create_joke():
 
 @pytest.mark.asyncio
 async def test_query_single_run():
+    # Enable tracing for this test - MUST be set before init()
+    os.environ["ALPHATRION_ENABLE_TRACING"] = "true"
+    os.environ["ALPHATRION_CLICKHOUSE_INIT_TABLES"] = "true"
+
     team_id = uuid.uuid4()
     user_id = uuid.uuid4()
 
@@ -413,7 +417,13 @@ async def test_query_single_run():
             exp_id = exp.id
             await exp.wait()
 
-    # Wait a bit for spans to be flushed to ClickHouse
+    # Force flush spans to ClickHouse
+    from opentelemetry import trace
+    tracer_provider = trace.get_tracer_provider()
+    if hasattr(tracer_provider, 'force_flush'):
+        tracer_provider.force_flush()
+
+    # Wait a bit for spans to be written to ClickHouse
     await asyncio.sleep(0.5)
 
     response = schema.execute_sync(