refactor: pass in JWTs instead of convoluted temp dir, fix test

Author: socksy

crates/tower-cmd/src/mcp.rs

@@ -1,6 +1,7 @@
 use anyhow::Result;
 use clap::Command;
 use crate::{Config, api, deploy, run};
+use config::Session;
 use rmcp::{
     ErrorData as McpError, ServerHandler,
     handler::server::{tool::{Parameters, ToolRouter}},
@@ -125,7 +126,11 @@ pub struct TowerService {
 impl TowerService {
     pub fn new(config: Config) -> Self {
         Self {
-            config,
+            config: std::env::var("TOWER_JWT")
+                .ok()
+                .and_then(|token| Session::from_jwt(&token).ok())
+                .map(|session| config.clone().with_session(session))
+                .unwrap_or(config),
             tool_router: Self::tool_router(),
         }
     }

tests/integration/features/environment.py

@@ -2,7 +2,6 @@
 import subprocess
 import time
 import tempfile
-import json
 import socket
 from pathlib import Path
 
@@ -23,12 +22,11 @@ def before_scenario(context, scenario):
 
     # Set up environment
     test_env = os.environ.copy()
-    test_env["TOWER_RUN_TIMEOUT"] = "1"
+    test_env["TOWER_RUN_TIMEOUT"] = "3"
 
-    # Create mock config if tower_url is set
     if context.tower_url:
         test_env["TOWER_URL"] = context.tower_url
-        _setup_mock_config(test_env, context.tower_url)
+        test_env["TOWER_JWT"] = "mock_jwt_token"
 
     # Find a free port for this test scenario
     mcp_port = _find_free_port()
@@ -38,7 +36,7 @@ def before_scenario(context, scenario):
         [tower_binary, "mcp-server", "--port", str(mcp_port)],
         env=test_env,
         stdout=subprocess.DEVNULL,
-        stderr=subprocess.PIPE,
+        stderr=subprocess.DEVNULL,
         text=True
     )
 
@@ -81,23 +79,6 @@ def _find_free_port():
         port = s.getsockname()[1]
     return port
 
-def _setup_mock_config(test_env, tower_url):
-    """Create a temporary tower configuration with mock session data"""
-    temp_config_dir = tempfile.mkdtemp(prefix="tower_test_config_")
-    test_env["HOME"] = temp_config_dir
-    
-    config_dir = os.path.join(temp_config_dir, ".config", "tower")
-    os.makedirs(config_dir, exist_ok=True)
-
-    mock_session = {
-        "user": {"id": "mock_user_id", "email": "test@example.com"},
-        "teams": [{"name": "default", "type": "user", "token": {"jwt": "mock_jwt_token"}}],
-        "active_team": {"name": "default", "type": "user", "token": {"jwt": "mock_jwt_token"}},
-        "tower_url": tower_url
-    }
-
-    with open(os.path.join(config_dir, "session.json"), 'w') as f:
-        json.dump(mock_session, f)
 
 def _find_tower_binary():
     # Look for debug build first

tests/integration/features/steps/mcp_steps.py

@@ -199,7 +199,7 @@ def step_check_timeout_message(context):
     assert hasattr(context, 'mcp_response'), "No MCP response was recorded"
 
     response_text = str(context.mcp_response).lower()
-    timeout_keywords = ["timeout", "timed out", "1 seconds"]
+    timeout_keywords = ["timeout", "timed out", "3 seconds"]
 
     found_timeout = any(keyword in response_text for keyword in timeout_keywords)
     assert found_timeout, f"Response should indicate timeout, got: {context.mcp_response}"

tests/integration/templates/long_runner.py

@@ -1,4 +1,4 @@
 import time
 print("Starting guaranteed-slow script (will timeout)...")
-time.sleep(10)
+time.sleep(5)
 print("This should never print")

tests/mock-api-server/main.py

@@ -190,6 +190,54 @@ async def describe_secrets_key():
 -----END RSA PUBLIC KEY-----"""
     return {"public_key": mock_public_key}
 
+def empty_paginated_response(key: str):
+    """Create empty paginated response for any resource type."""
+    return {
+        key: [],
+        "pages": {"page": 1, "total": 0, "num_pages": 1, "page_size": 20}
+    }
+
+@app.post("/v1/secrets/export")
+async def export_secrets(export_params: Dict[str, Any]):
+    """Mock endpoint for exporting secrets with encryption."""
+    return empty_paginated_response("secrets")
+
+@app.post("/v1/catalogs/export")
+async def export_catalogs(export_params: Dict[str, Any]):
+    """Mock endpoint for exporting catalogs with encryption."""
+    return empty_paginated_response("catalogs")
+
+@app.get("/v1/session")
+async def get_session():
+    """Mock endpoint for getting current session."""
+    return {
+        "session": {
+            "featurebase_identity": {
+                "company_hash": "mock_company_hash",
+                "user_hash": "mock_user_hash"
+            },
+            "user": {
+                "company": "Mock Company",
+                "country": "US",
+                "created_at": "2023-01-01T00:00:00Z",
+                "email": "test@example.com",
+                "first_name": "Test",
+                "is_alerts_enabled": True,
+                "is_invitation_claimed": True,
+                "last_name": "User",
+                "profile_photo_url": "https://example.com/photo.jpg"
+            },
+            "teams": [
+                {
+                    "name": "default",
+                    "type": "user",
+                    "token": {"jwt": "mock_jwt_token"}
+                }
+            ],
+            "token": {"jwt": "mock_jwt_token"}
+        }
+    }
+
 @app.post("/v1/session/refresh")
 async def refresh_session(refresh_params: Dict[str, Any] = None):
     """Mock endpoint for refreshing session."""