From 9189b524538e0a0a7761d31e04f094075cf5659e Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Thu, 19 Dec 2024 21:29:38 +0200
Subject: [PATCH] fix(add tests):

---
 Readme.md                     |  2 ++
 agentic_security/http_spec.py |  3 +-
 agentic_security/test_spec.py | 64 +++++++++++++++++++++++++++++++++++
 poetry.lock                   | 20 ++++++++++-
 pyproject.toml                |  1 +
 5 files changed, 88 insertions(+), 2 deletions(-)

diff --git a/Readme.md b/Readme.md
index b7f79d7..934fcfa 100644
--- a/Readme.md
+++ b/Readme.md
@@ -286,6 +286,7 @@ Content-Type: application/json
 ```
 
 Replace `XXXXX` with your actual API key and `<<BASE64_IMAGE>>` is the image variable.
+
 ## Audio Modality
 
 To probe the audio modality, you can use the following HTTP request:
@@ -302,6 +303,7 @@ Content-Type: multipart/form-data
 ```
 
 Replace `$GROQ_API_KEY` with your actual API key and ensure that the `file` parameter points to the correct audio file path.
+
 ## CI/CD integration
 
 TBD
diff --git a/agentic_security/http_spec.py b/agentic_security/http_spec.py
index 2a52507..9471a52 100644
--- a/agentic_security/http_spec.py
+++ b/agentic_security/http_spec.py
@@ -2,7 +2,8 @@ import httpx
 from pydantic import BaseModel
 
 
-class InvalidHTTPSpecError(Exception): ...
+class InvalidHTTPSpecError(Exception):
+    ...
 
 
 class LLMSpec(BaseModel):
diff --git a/agentic_security/test_spec.py b/agentic_security/test_spec.py
index 46c9687..fb31a45 100644
--- a/agentic_security/test_spec.py
+++ b/agentic_security/test_spec.py
@@ -1,3 +1,5 @@
+import pytest
+
 from agentic_security.http_spec import LLMSpec, parse_http_spec
 
 
@@ -52,3 +54,65 @@ class TestParseHttpSpec:
         assert result.url == "http://example.com"
         assert result.headers == {"Content-Type": "application/json"}
         assert result.body == ""
+
+
+class TestLLMSpec:
+    def test_validate_raises_error_for_missing_files(self):
+        spec = LLMSpec(
+            method="POST", url="http://example.com", headers={}, body="", has_files=True
+        )
+        with pytest.raises(ValueError, match="Files are required for this request."):
+            spec.validate(prompt="", encoded_image="", files={})
+
+    def test_validate_raises_error_for_missing_image(self):
+        spec = LLMSpec(
+            method="POST", url="http://example.com", headers={}, body="", has_image=True
+        )
+        with pytest.raises(ValueError, match="An image is required for this request."):
+            spec.validate(prompt="", encoded_image="", files={})
+
+    @pytest.mark.asyncio
+    async def test_probe_sends_request(self, httpx_mock):
+        httpx_mock.add_response(
+            method="POST", url="http://example.com", status_code=200
+        )
+        spec = LLMSpec(
+            method="POST",
+            url="http://example.com",
+            headers={},
+            body='{"prompt": "<<PROMPT>>"}',
+        )
+        response = await spec.probe(prompt="test")
+        assert response.status_code == 200
+
+    @pytest.mark.asyncio
+    async def test_probe_with_files(self, httpx_mock):
+        httpx_mock.add_response(
+            method="POST", url="http://example.com", status_code=200
+        )
+        spec = LLMSpec(
+            method="POST",
+            url="http://example.com",
+            headers={"Content-Type": "multipart/form-data"},
+            body='{"prompt": "<<PROMPT>>"}',
+            has_files=True,
+        )
+        files = {"file": ("filename.txt", "file content")}
+        response = await spec.probe(prompt="test", files=files)
+        assert response.status_code == 200
+
+    @pytest.mark.asyncio
+    async def test_probe_with_image(self, httpx_mock):
+        httpx_mock.add_response(
+            method="POST", url="http://example.com", status_code=200
+        )
+        spec = LLMSpec(
+            method="POST",
+            url="http://example.com",
+            headers={},
+            body='{"image": "<<BASE64_IMAGE>>"}',
+            has_image=True,
+        )
+        encoded_image = "base64encodedstring"
+        response = await spec.probe(prompt="test", encoded_image=encoded_image)
+        assert response.status_code == 200
diff --git a/poetry.lock b/poetry.lock
index 121f6b5..8dd629a 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -2335,6 +2335,24 @@ pytest = ">=8.2,<9"
 docs = ["sphinx (>=5.3)", "sphinx-rtd-theme (>=1.0)"]
 testing = ["coverage (>=6.2)", "hypothesis (>=5.7.1)"]
 
+[[package]]
+name = "pytest-httpx"
+version = "0.35.0"
+description = "Send responses to httpx."
+optional = false
+python-versions = ">=3.9"
+files = [
+    {file = "pytest_httpx-0.35.0-py3-none-any.whl", hash = "sha256:ee11a00ffcea94a5cbff47af2114d34c5b231c326902458deed73f9c459fd744"},
+    {file = "pytest_httpx-0.35.0.tar.gz", hash = "sha256:d619ad5d2e67734abfbb224c3d9025d64795d4b8711116b1a13f72a251ae511f"},
+]
+
+[package.dependencies]
+httpx = "==0.28.*"
+pytest = "==8.*"
+
+[package.extras]
+testing = ["pytest-asyncio (==0.24.*)", "pytest-cov (==6.*)"]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"
@@ -3011,4 +3029,4 @@ propcache = ">=0.2.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "24d9cc0ad18be478dd235d69f920e78d0daabe64a3286186f6f79c3990dcc41b"
+content-hash = "df0a6b4d6fe2aea83a5e4b59ed9c4913a7cdf45ac0a7e668d4c3740c7197340c"
diff --git a/pyproject.toml b/pyproject.toml
index ffff05d..a975e00 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -53,6 +53,7 @@ pytest = "^8.3.4"
 pre-commit = "^4.0.1"
 langchain-groq = "^0.2.0"
 huggingface-hub = "^0.25.1"
+pytest-httpx = "^0.35.0"
 # garak = "*"
 
 [tool.ruff]