test: add test for encryption

Author: WHALEEYE

crab/server/api.py

@@ -26,17 +26,17 @@
 from .logger import crab_logger as logger
 
 api_router = APIRouter()
-ENC_KEY = generate_key_from_env()
 
 
 @api_router.post("/raw_action")
 async def raw_action(request: Request):
     """Perform the specified action with given parameters."""
+    enc_key = generate_key_from_env()
     # Extract query parameters as a dictionary
     request_content = await request.body()
     request_content = request_content.decode("utf-8")
-    if ENC_KEY is not None:
-        request_content = decrypt_message(request_content, ENC_KEY)
+    if enc_key is not None:
+        request_content = decrypt_message(request_content, enc_key)
     request_json = json.loads(request_content)
 
     action = request_json["action"]
@@ -47,8 +47,8 @@ async def raw_action(request: Request):
         parameters["env"] = request.app.environment
 
     resp_data = {"action_returns": entry(**parameters)}
-    if ENC_KEY is None:
+    if enc_key is None:
         return JSONResponse(content=resp_data)
     else:
-        encrypted = encrypt_message(json.dumps(resp_data), ENC_KEY)
+        encrypted = encrypt_message(json.dumps(resp_data), enc_key)
         return PlainTextResponse(content=encrypted)

pyproject.toml

@@ -122,7 +122,7 @@ certifi = "^2024.2.2"
 lint.select = ["E501", "E4", "E7", "E9", "F", "I"]
 lint.ignore = ["E731"]
 lint.ignore-init-module-imports = true
-exclude = ["docs/"]
+exclude = ["docs/", "crab-benchmark-v0/thirdparty/"]
 
 [[tool.mypy.overrides]]
 module = ["dill", "easyocr"]

test/server/test_api.py

@@ -11,8 +11,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # =========== Copyright 2024 @ CAMEL-AI.org. All Rights Reserved. ===========
+
+import pytest
 from fastapi.testclient import TestClient
 
+from crab import create_environment
 from crab.environments.template import (
     current_state,
     set_state,
@@ -21,53 +24,27 @@
 from crab.server.main import init
 
 
-def test_raw_action():
-    app = init(template_environment_config)
-    client = TestClient(app)
-    response = client.post(
-        "/raw_action",
-        json={
-            "action": set_state.to_raw_action(),
-            "parameters": {"value": True},
-        },
-    )
-    assert response.json()["action_returns"] is None
+@pytest.fixture
+def mock_env():
+    mock_app = init(template_environment_config)
+    mock_cli = TestClient(mock_app)
+    mock_env = create_environment(template_environment_config)
+    mock_env._client = mock_cli
+    return mock_env
 
-    response = client.post(
-        "/raw_action",
-        json={
-            "action": current_state.to_raw_action(),
-            "parameters": {},
-        },
-    )
-    assert response.json()["action_returns"] is True
 
-    action = set_state(True)
-    response = client.post(
-        "/raw_action",
-        json={
-            "action": action.to_raw_action(),
-            "parameters": {},
-        },
-    )
-    assert response.json()["action_returns"] is None
+def test_raw_action_unencrypted(mock_env):
+    assert mock_env._action_endpoint(set_state, {"value": True}) is None
+    assert mock_env._action_endpoint(current_state, {}) is True
+    assert mock_env._action_endpoint(set_state(True), {}) is None
+    assert mock_env._action_endpoint(current_state >> set_state, {}) is None
+    assert mock_env._action_endpoint(set_state(True) + current_state, {}) is True
 
-    action = current_state >> set_state
-    response = client.post(
-        "/raw_action",
-        json={
-            "action": action.to_raw_action(),
-            "parameters": {},
-        },
-    )
-    assert response.json()["action_returns"] is None
 
-    action = set_state(True) + current_state
-    response = client.post(
-        "/raw_action",
-        json={
-            "action": action.to_raw_action(),
-            "parameters": {},
-        },
-    )
-    assert response.json()["action_returns"] is True
+def test_raw_action_encrypted(mock_env, monkeypatch):
+    monkeypatch.setenv("ENCRYPTION_KEY", "the-cake-is-a-lie")
+    assert mock_env._action_endpoint(set_state, {"value": True}) is None
+    assert mock_env._action_endpoint(current_state, {}) is True
+    assert mock_env._action_endpoint(set_state(True), {}) is None
+    assert mock_env._action_endpoint(current_state >> set_state, {}) is None
+    assert mock_env._action_endpoint(set_state(True) + current_state, {}) is True