"Add performance testing function and update atomic.py"

backend/api/tests/atomic.py

@@ -2,57 +2,76 @@
 import time
 import atomic_bomb_engine
 from atomic_bomb_engine import server
-from unitrunner.engine.bomb import AtomicBombEngine
+from unitrunner.engine.bomb import AtomicBombEngine, performance
 
 
-@server.ui(port=8880)
-async def main():
-    engine = AtomicBombEngine(
+# @server.ui(port=8880)
+# async def main():
+#     engine = AtomicBombEngine(
+#         test_duration_secs=10,
+#         concurrent_requests=1,
+#         timeout_secs=15,
+#         cookie_store_enable=True,
+#         verbose=False,
+#         increase_step=1,
+#         increase_interval=3
+#     )
+#     validators = [{
+#         "jsonpath": "$.msg",
+#         "reference_object": "操作成功"
+#     }]
+#     runner = engine.run(data={
+#         "name": "test-1",
+#         "url": "https://www.baidu.com/",
+#         "method": "POST",
+#         "form_data": {"name": "{{api-test-msg-1}}", "number": "{{api-test-num}}"},
+#         "validators": validators,
+#     })
+#     # runner = atomic_bomb_engine.BatchRunner()
+#     # step = getattr(atomic_bomb_engine, 'endpoint')
+#     # runner.run(
+#     #     test_duration_secs=10,
+#     #     concurrent_requests=1,
+#     #     step_option=atomic_bomb_engine.step_option(increase_step=1, increase_interval=3),
+#     #     timeout_secs=15,
+#     #     cookie_store_enable=True,
+#     #     verbose=False,
+#     #     api_endpoints=[
+#     #         step(
+#     #             name="test-1",
+#     #             url="https://www.baidu.com/",
+#     #             method="POST",
+#     #             form_data={"name": "{{api-test-msg-1}}", "number": "{{api-test-num}}"},
+#     #             weight=100,
+#     #             assert_options=[
+#     #                 atomic_bomb_engine.assert_option(jsonpath="$.msg", reference_object="操作成功"),
+#     #             ],
+#     #         ),
+#     #     ])
+#     return runner
+
+
+if __name__ == '__main__':
+    data = {
+        "name": "test-1",
+        "url": "https://www.baidu.com/",
+        "method": "POST",
+        "form_data": {"name": "{{api-test-msg-1}}", "number": "{{api-test-num}}"},
+        "validators": [{
+            "jsonpath": "$.msg",
+            "reference_object": "操作成功"
+        }],
+    }
+    results = asyncio.run(performance(
+        test_data=data,
         test_duration_secs=10,
         concurrent_requests=1,
         timeout_secs=15,
         cookie_store_enable=True,
         verbose=False,
         increase_step=1,
         increase_interval=3
-    )
-    validators = [{
-        "jsonpath": "$.msg",
-        "reference_object": "操作成功"
-    }]
-    runner = engine.run(data={
-        "name": "test-1",
-        "url": "https://www.baidu.com/",
-        "method": "POST",
-        "form_data": {"name": "{{api-test-msg-1}}", "number": "{{api-test-num}}"},
-        "validators": validators,
-    })
-    # runner = atomic_bomb_engine.BatchRunner()
-    # step = getattr(atomic_bomb_engine, 'endpoint')
-    # runner.run(
-    #     test_duration_secs=10,
-    #     concurrent_requests=1,
-    #     step_option=atomic_bomb_engine.step_option(increase_step=1, increase_interval=3),
-    #     timeout_secs=15,
-    #     cookie_store_enable=True,
-    #     verbose=False,
-    #     api_endpoints=[
-    #         step(
-    #             name="test-1",
-    #             url="https://www.baidu.com/",
-    #             method="POST",
-    #             form_data={"name": "{{api-test-msg-1}}", "number": "{{api-test-num}}"},
-    #             weight=100,
-    #             assert_options=[
-    #                 atomic_bomb_engine.assert_option(jsonpath="$.msg", reference_object="操作成功"),
-    #             ],
-    #         ),
-    #     ])
-    return runner
-
-
-if __name__ == '__main__':
-    results = asyncio.run(main())
+    ))
     a = []
     for res in results:
         if res.get("should_wait"):

backend/unitrunner/engine/bomb.py

@@ -167,3 +167,45 @@ def assert_option(validate_check: list) -> List[Any]:
             assert_options.append(atomic_bomb_engine.assert_option(jsonpath, reference_object))
 
         return assert_options
+
+
+async def performance(
+        test_data,
+        test_duration_secs=None,
+        concurrent_requests=None,
+        timeout_secs=0,
+        cookie_store_enable=True,
+        verbose=False,
+        increase_step=None,
+        increase_interval=None
+):
+    """
+    Runs a performance test with the provided configuration.
+
+    Args:
+        test_data: Data to be used for the test.
+        test_duration_secs: Optional duration of the test in seconds.
+        concurrent_requests: Optional number of concurrent requests.
+        timeout_secs: Optional timeout in seconds for individual requests.
+        cookie_store_enable: Optional flag to enable cookie storage.
+        verbose: Optional flag to enable verbose logging.
+        increase_step: Optional step size for increasing load (experimental).
+        increase_interval: Optional interval for load increase (experimental).
+
+    Returns:
+        An object representing the test runner (implementation-specific).
+    """
+    try:
+        engine = AtomicBombEngine(
+            test_duration_secs=test_duration_secs,
+            concurrent_requests=concurrent_requests,
+            timeout_secs=timeout_secs,
+            cookie_store_enable=cookie_store_enable,
+            verbose=verbose,
+            increase_step=increase_step,
+            increase_interval=increase_interval,
+        )
+        runner = engine.run(test_data)
+        return runner
+    except (Exception, KeyboardInterrupt):
+        raise Exception("Test runner encountered an exception.")