updated test cases

Signed-off-by: Keval Mahajan <[email protected]>

Author: kevalmahajan

tests/unit/mcpgateway/services/test_gateway_service.py

@@ -570,17 +570,45 @@ async def test_register_gateway_with_existing_tools(self, gateway_service, test_
             (200, {"content-type": "application/json"}, "SSE", False),
         ],
     )
+    
+    # helper to make SSE validation work:
+    @pytest.mark.parametrize(
+        "status_code,headers,transport_type,expected",
+        [
+            # SSE transport success cases
+            (200, {"content-type": "text/event-stream"}, "SSE", True),
+            # SSE transport failure cases - auth failures
+            (401, {"content-type": "text/event-stream"}, "SSE", False),
+            (403, {"content-type": "text/event-stream"}, "SSE", False),
+            # SSE transport failure cases - wrong content-type
+            (200, {"content-type": "application/json"}, "SSE", False),
+        ],
+    )
     @pytest.mark.asyncio
     async def test_validate_gateway_url_responses(self, gateway_service, httpx_mock, status_code, headers, transport_type, expected):
         """Test various HTTP responses during gateway URL validation."""
-        httpx_mock.add_response(
-            method="GET",
-            url="http://example.com",
-            status_code=status_code,
-            headers=headers,
-        )
+        method = "POST" if transport_type == "STREAMABLEHTTP" else "GET"
+        
+        # For SSE with 200 status, mock streaming response
+        if transport_type == "SSE" and status_code == 200 and "text/event-stream" in headers.get("content-type", ""):
+            httpx_mock.add_response(
+                method=method,
+                url="http://example.com",
+                status_code=status_code,
+                headers=headers,
+                content=b"data: test\n\n",  # Add SSE data so aiter_lines() returns something
+            )
+        else:
+            httpx_mock.add_response(
+                method=method,
+                url="http://example.com",
+                status_code=status_code,
+                headers=headers,
+            )
 
-        result = await gateway_service._validate_gateway_url(url="http://example.com", headers={}, transport_type=transport_type)
+        result = await gateway_service._validate_gateway_url(
+            url="http://example.com", headers={}, transport_type=transport_type
+        )
 
         assert result is expected
 
@@ -617,25 +645,19 @@ async def test_ssl_verification_bypass(self, gateway_service, monkeypatch):
     @pytest.mark.asyncio
     async def test_streamablehttp_redirect(self, gateway_service, httpx_mock):
         """Test STREAMABLEHTTP transport with redirection and MCP session ID."""
-        # Mock first response with redirect
+        # When follow_redirects=True, httpx handles redirects internally
+        # Only mock the FINAL response, not intermediate redirects
         httpx_mock.add_response(
-            method="GET",
+            method="POST",
             url="http://example.com",
-            status_code=302,
-            headers={"location": "http://sampleredirected.com"},
-        )
-
-        # Mock redirected response with MCP session
-        httpx_mock.add_response(
-            method="GET",
-            url="http://sampleredirected.com",
             status_code=200,
-            headers={"mcp-session-id": "sample123", "content-type": "application/json"},
+            headers={"content-type": "application/json"},
         )
 
-        result = await gateway_service._validate_gateway_url(url="http://example.com", headers={}, transport_type="STREAMABLEHTTP")
+        result = await gateway_service._validate_gateway_url(
+            url="http://example.com", headers={}, transport_type="STREAMABLEHTTP"
+        )
 
-        # Should return True when redirect has mcp-session-id and application/json content-type
         assert result is True
 
     # ───────────────────────────────────────────────────────────────────────────
@@ -645,22 +667,23 @@ async def test_streamablehttp_redirect(self, gateway_service, httpx_mock):
     async def test_bulk_concurrent_validation(self, gateway_service, httpx_mock):
         """Test bulk concurrent gateway URL validations."""
         urls = [f"http://gateway{i}.com" for i in range(20)]
-
-        # Add responses for all URLs
+        
+        # Add responses for all URLs with SSE content
         for url in urls:
             httpx_mock.add_response(
                 method="GET",
                 url=url,
                 status_code=200,
                 headers={"content-type": "text/event-stream"},
+                content=b"data: test\n\n",  # Add SSE data
             )
 
         # Run the validations concurrently
         results = await asyncio.gather(*[gateway_service._validate_gateway_url(url, {}, "SSE") for url in urls])
 
         # All should be True (validation success)
         assert all(results)
-
+        
     # ────────────────────────────────────────────────────────────────────
     # LIST / GET
     # ────────────────────────────────────────────────────────────────────
@@ -1322,47 +1345,34 @@ async def test_forward_request_connection_error(self, gateway_service, mock_gate
     @pytest.mark.asyncio
     async def test_validate_gateway_url_redirect_with_auth_failure(self, gateway_service, httpx_mock):
         """Test redirect handling with authentication failure at redirect location."""
-        # Mock first response (redirect with Location header)
+        # Only mock final response with auth failure
         httpx_mock.add_response(
-            method="GET",
+            method="POST",
             url="http://example.com",
-            status_code=302,
-            headers={"location": "http://redirected.com/api"},
-        )
-
-        # Mock redirected response with auth failure
-        httpx_mock.add_response(
-            method="GET",
-            url="http://redirected.com/api",
             status_code=401,
         )
 
-        result = await gateway_service._validate_gateway_url(url="http://example.com", headers={}, transport_type="STREAMABLEHTTP")
+        result = await gateway_service._validate_gateway_url(
+            url="http://example.com", headers={}, transport_type="STREAMABLEHTTP"
+        )
 
         assert result is False
 
     @pytest.mark.asyncio
     async def test_validate_gateway_url_redirect_with_mcp_session(self, gateway_service, httpx_mock):
         """Test redirect handling with MCP session ID in response."""
-        # Mock first response (redirect with Location header)
+        # STREAMABLEHTTP uses POST method, and only mock final response
         httpx_mock.add_response(
-            method="GET",
+            method="POST",  # Changed from GET to POST
             url="http://example.com",
-            status_code=302,
-            headers={"location": "http://redirected.com/api"},
-        )
-
-        # Mock redirected response with MCP session
-        httpx_mock.add_response(
-            method="GET",
-            url="http://redirected.com/api",
             status_code=200,
             headers={"mcp-session-id": "session123", "content-type": "application/json"},
         )
 
-        result = await gateway_service._validate_gateway_url(url="http://example.com", headers={}, transport_type="STREAMABLEHTTP")
+        result = await gateway_service._validate_gateway_url(
+            url="http://example.com", headers={}, transport_type="STREAMABLEHTTP"
+        )
 
-        # Should return True when redirect has mcp-session-id and application/json content-type
         assert result is True
 
     # ────────────────────────────────────────────────────────────────────