From bd2dbfefd7ff55f2cf20042b064fdfe1c51ffcf1 Mon Sep 17 00:00:00 2001
From: "sweep-ai[bot]" <128439645+sweep-ai[bot]@users.noreply.github.com>
Date: Thu, 9 Nov 2023 16:54:14 +0000
Subject: [PATCH] feat: Add unit tests for subprocessor module

---
 test/test_subprocessor.py | 51 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)
 create mode 100644 test/test_subprocessor.py

diff --git a/test/test_subprocessor.py b/test/test_subprocessor.py
new file mode 100644
index 0000000..e711768
--- /dev/null
+++ b/test/test_subprocessor.py
@@ -0,0 +1,51 @@
+import unittest
+from unittest.mock import patch, call
+from pulling_ace.utils import subprocessor
+
+class SubprocessorTest(unittest.TestCase):
+    @patch('subprocessor.subprocess.run')
+    def test_toxicity(self, mock_run):
+        subprocessor.toxicity('huggingface', 'gpt2', 'RTPBlank')
+        mock_run.assert_called_once_with(
+            ['python3', '-m', 'garak', '--model_type', 'huggingface', '--model_name', 'gpt2', '--probes', 'realtoxicityprompts.RTPBlank'],
+            check=True, capture_output=True, text=True
+        )
+
+    @patch('subprocessor.subprocess.run')
+    def test_promptinjection(self, mock_run):
+        subprocessor.promptinjection('huggingface', 'gpt2', 'HijackHateHumans')
+        mock_run.assert_called_once_with(
+            ['python3', '-m', 'garak', '--model_type', 'huggingface', '--model_name', 'gpt2', '--probes', 'promptinject.HijackHateHumans'],
+            check=True, capture_output=True, text=True
+        )
+
+    @patch('subprocessor.subprocess.run')
+    def test_riskcards(self, mock_run):
+        subprocessor.riskcards('huggingface', 'gpt2', 'Bullying')
+        mock_run.assert_called_once_with(
+            ['python3', '-m', 'garak', '--model_type', 'huggingface', '--model_name', 'gpt2', '--probes', 'lmrc.Bullying'],
+            check=True, capture_output=True, text=True
+        )
+
+    @patch('subprocessor.riskcards')
+    def test_riskcard_wrapper(self, mock_riskcards):
+        subprocessor.riskcard_wrapper({'model_type': 'huggingface', 'model_name': 'gpt2', 'probe': 'Bullying'})
+        mock_riskcards.assert_called_once_with('huggingface', 'gpt2', 'Bullying')
+
+    @patch('subprocessor.promptinjection')
+    def test_promptinjection_wrapper(self, mock_promptinjection):
+        subprocessor.promptinjection_wrapper({'model_type': 'huggingface', 'model_name': 'gpt2', 'probe': 'HijackHateHumans'})
+        mock_promptinjection.assert_called_once_with('huggingface', 'gpt2', 'HijackHateHumans')
+
+    @patch('subprocessor.toxicity')
+    def test_toxicity_wrapper(self, mock_toxicity):
+        subprocessor.toxicity_wrapper({'model_type': 'huggingface', 'model_name': 'gpt2', 'probe': 'RTPBlank'})
+        mock_toxicity.assert_called_once_with('huggingface', 'gpt2', 'RTPBlank')
+
+    @patch('subprocessor.multiprocessing.Pool')
+    def test_run_injections(self, mock_pool):
+        subprocessor.run_injections('huggingface', 'gpt2', 'promptinject')
+        mock_pool.return_value.map.assert_called_once()
+
+if __name__ == '__main__':
+    unittest.main()