from fastapi.testclient import TestClient from main import app # Replace with the actual name of your FastAPI app import json client = TestClient(app) def test_open_code_in_vs_code_valid(): response = client.post( "/open_code_in_vs_code", data=json.dumps({ "discussion_id": 1, "message_id": 1, "code": "print('Hello, World!')" }), headers={"content-type": "application/json"}, ) assert response.status_code == 200 assert response.json()["status"] == True def test_open_code_in_vs_code_invalid(): response = client.post( "/open_code_in_vs_code", data=json.dumps({ "discussion_id": "1; copy file.exe /some/path/", "message_id": "1", "code": "print('Hello, World!')" }), headers={"content-type": "application/json"}, ) assert response.status_code == 422 # Unprocessable Entity def test_open_code_in_vs_code_attack(): response = client.post( "/open_code_in_vs_code", data=json.dumps({ "discussion_id": 1, "message_id": 1, "code": "print('This is a harmless test.')" # Dangerous code }), headers={"content-type": "application/json"}, ) assert response.status_code == 200 assert response.json()["status"] == False # The code should not be executed