From ad9b810083045da1009df61fa5d77a39fb236463 Mon Sep 17 00:00:00 2001 From: dmeenaarmorcode Date: Fri, 31 Jan 2025 09:25:53 +0530 Subject: [PATCH 1/3] Test Demo Pr --- web-agent/Test/AgentIntegrationTest.py | 124 ++++++++++++++ web-agent/Test/HttpTeleportTask.py | 5 + web-agent/Test/LocalTest.py | 36 +++++ web-agent/Test/MockClientServer.py | 46 ++++++ web-agent/Test/Mock_ArmorCode_server.py | 153 ++++++++++++++++++ .../Test/Resources/Files/json_response.json | 0 web-agent/Test/Resources/TestCases.py | 59 +++++++ web-agent/Test/TaskStatus.py | 7 + web-agent/Test/VerifyResult.py | 48 ++++++ web-agent/Test/config.properties | 1 + web-agent/app/worker.py | 4 +- 11 files changed, 481 insertions(+), 2 deletions(-) create mode 100644 web-agent/Test/AgentIntegrationTest.py create mode 100644 web-agent/Test/HttpTeleportTask.py create mode 100644 web-agent/Test/LocalTest.py create mode 100644 web-agent/Test/MockClientServer.py create mode 100644 web-agent/Test/Mock_ArmorCode_server.py create mode 100644 web-agent/Test/Resources/Files/json_response.json create mode 100644 web-agent/Test/Resources/TestCases.py create mode 100644 web-agent/Test/TaskStatus.py create mode 100644 web-agent/Test/VerifyResult.py create mode 100644 web-agent/Test/config.properties diff --git a/web-agent/Test/AgentIntegrationTest.py b/web-agent/Test/AgentIntegrationTest.py new file mode 100644 index 0000000..da19df9 --- /dev/null +++ b/web-agent/Test/AgentIntegrationTest.py @@ -0,0 +1,124 @@ +import os +import time +import unittest +import subprocess +import threading +import zipfile +from random import random +from urllib.request import urlretrieve + +import Mock_ArmorCode_server +from MockClientServer import client_app +from Resources import TestCases + +file_directory = "Resources/Files/FilesToUpload" + +class TestArmorCodeAgent(unittest.TestCase): + @classmethod + def setUpClass(cls): + # Create necessary directories + os.makedirs(file_directory, exist_ok=True) + + # Create files mentioned in TestCases.py + # for case in TestCases.test_cases: + # if case['response_file_path']: + # file_path = os.path.join(file_directory, case['response_file_path']) + # if case['task_id'] == 'big_file_task': + # # Create a 300 MB zip file + # cls.create_large_zip_file(file_path, 300 * 1024 * 1024) # 300 MB in bytes + # else: + # with open(file_path, 'w') as f: + # f.write(f"Mock content for {case['task_id']}") + + # Download worker.py and requirements.txt + urlretrieve('https://raw.githubusercontent.com/armor-code/agent/refs/heads/main/web-agent/app/worker.py', 'worker.py') + urlretrieve('https://raw.githubusercontent.com/armor-code/agent/refs/heads/main/web-agent/requirements.txt', 'requirements.txt') + + # Install requirements + subprocess.run(['pip3', 'install', '-r', 'requirements.txt']) + + # Start Mock_ArmorCode_server + # cls.ac_thread = threading.Thread(target=lambda: ac_app.run(port=5000)) + # cls.ac_thread.daemon = True + # cls.ac_thread.start() + # + # # Start MockClientServer + # cls.client_thread = threading.Thread(target=lambda: client_app.run(port=5001)) + # cls.client_thread.daemon = True + # cls.client_thread.start() + # + # # Wait for servers to start + # time.sleep(2) + + @staticmethod + def create_large_zip_file(file_path, size_bytes): + with zipfile.ZipFile(file_path, 'w', zipfile.ZIP_DEFLATED) as zip_file: + # Create a large file with random data + chunk_size = 1024 * 1024 # 1 MB + remaining_size = size_bytes + while remaining_size > 0: + chunk = os.urandom(min(chunk_size, remaining_size)) + + zip_file.writestr(f"data_{remaining_size}.bin", chunk) + remaining_size -= len(chunk) + + def setUp(self): + # Clear task_result_map before each test + Mock_ArmorCode_server.task_result_map.clear() + + def run_worker_and_wait(self, task_id): + # Start worker process + worker_process = subprocess.Popen(['python3', 'worker.py', '--serverUrl', 'http://localhost:5000', '--apiKey', 'test_api_key', '--index', 'test_index', '--timeout', '30', '--verify', 'False', '--debugMode', 'False']); + + # Wait for task result (max 60 seconds) + start_time = time.time() + while time.time() - start_time < 60: + if task_id in Mock_ArmorCode_server.task_result_map: + break + time.sleep(1) + + # Stop worker process + worker_process.terminate() + worker_process.wait() + + return Mock_ArmorCode_server.task_result_map.get(task_id) + + def test_tasks(self): + for case in TestCases.test_cases: + with self.subTest(task_id=case['task_id']): + # Set task type to return + Mock_ArmorCode_server.write_global_var(case['task_id']) + + # Run worker and wait for result + result = self.run_worker_and_wait(case['task_id']) + + # Assert that we got a result + self.assertIsNotNone(result, f"No result received for task {case['task_id']}") + + # Add more specific assertions based on expected results + # For example: + if case['task_id'] == 'json_task': + self.assertEqual(result['status'], 'success') + self.assertIn('data', result) + elif case['task_id'] == 'timeout_task': + self.assertEqual(result['status'], 'error') + self.assertIn('timeout', result['message'].lower()) + # Add more assertions for other task types + + @classmethod + def tearDownClass(cls): + # Clean up created files + for case in TestCases.test_cases: + if case['response_file_path']: + file_path = os.path.join('Test/Resources/Files/S3_upload', case['response_file_path']) + if os.path.exists(file_path): + os.remove(file_path) + + # Remove downloaded files + if os.path.exists('worker.py'): + os.remove('worker.py') + if os.path.exists('requirements.txt'): + os.remove('requirements.txt') + +if __name__ == '__main__': + unittest.main() \ No newline at end of file diff --git a/web-agent/Test/HttpTeleportTask.py b/web-agent/Test/HttpTeleportTask.py new file mode 100644 index 0000000..e05af42 --- /dev/null +++ b/web-agent/Test/HttpTeleportTask.py @@ -0,0 +1,5 @@ +class HttpTeleportTask: + def __init__(self, task_id, url, method): + self.task_id = task_id + self.url = url + self.method = method diff --git a/web-agent/Test/LocalTest.py b/web-agent/Test/LocalTest.py new file mode 100644 index 0000000..993a03b --- /dev/null +++ b/web-agent/Test/LocalTest.py @@ -0,0 +1,36 @@ +import os + +import requests + + +def test_mock_s3_upload(): + # URL of the mock S3 upload endpoint + url = 'http://localhost:5000/mock-s3-upload' + + # Path to a test file + test_file_path = 'test_file.txt' + + # Create a test file + with open(test_file_path, 'w') as f: + f.write('This is a test file for mock S3 upload.') + + # Open the file and send it in a PUT request + with open(test_file_path, 'rb') as f: + files = {'file': ('test_file.txt', f)} + response = requests.put(url, files=files) + + # Print the response + print(f"Status Code: {response.status_code}") + print(f"Response: {response.json()}") + + # Clean up - remove the test file + os.remove(test_file_path) + + # Check if the file was uploaded to the S3_upload folder + uploaded_file_path = os.path.join('Resources', 'Files', 'S3_upload', 'test_file.txt') + if os.path.exists(uploaded_file_path): + print(f"File successfully uploaded to {uploaded_file_path}") + else: + print("File upload failed") + +test_mock_s3_upload() \ No newline at end of file diff --git a/web-agent/Test/MockClientServer.py b/web-agent/Test/MockClientServer.py new file mode 100644 index 0000000..99e378c --- /dev/null +++ b/web-agent/Test/MockClientServer.py @@ -0,0 +1,46 @@ +import os + +from flask import Flask, request, jsonify, send_file +from werkzeug.serving import run_simple +import Resources.TestCases as TestCases + +client_app = Flask("client_server") +test_cases = TestCases.test_cases +BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +UPLOAD_DIR = os.path.join(BASE_DIR, 'Test', 'Resources', 'Files', 'S3_upload') + + +client_app = Flask("client_server") + + +@client_app.before_request +def initialize_tasks(): + return + + +@client_app.route('/', methods=['GET']) +def serve_file(path): + for case in test_cases: + if case['url_path'] == f"/{path}": + if path == "timeout": + return "", 504 + elif path == "error_400": + return "", 400 + elif path == "error_500": + return "", 500 + elif case['response_file_path']: + # Construct the full path to the file + file_path = os.path.join(UPLOAD_DIR, os.path.basename(case['response_file_path'])) + if os.path.exists(file_path): + return send_file( + file_path, + mimetype=case['headers'].get('Content-Type', 'application/octet-stream'), + as_attachment=True + ) + else: + return f"File not found: {file_path}", 404 + return "", 404 + +if __name__ == "__main__": + # Run both servers + run_simple('localhost', 5001, client_app, use_reloader=True, use_debugger=True, use_evalex=True) \ No newline at end of file diff --git a/web-agent/Test/Mock_ArmorCode_server.py b/web-agent/Test/Mock_ArmorCode_server.py new file mode 100644 index 0000000..53fdfdd --- /dev/null +++ b/web-agent/Test/Mock_ArmorCode_server.py @@ -0,0 +1,153 @@ +import json +import gzip +import os + +import requests +from flask import Flask, request, jsonify, send_file +from werkzeug.serving import run_simple +import Resources.TestCases as TestCases + +from werkzeug.utils import secure_filename +from TaskStatus import TaskStatusEnum +from VerifyResult import VerificationResult, verify_task_result + +import HttpTeleportTask + +# Mock Armorcode Server +ac_app = Flask("ac_server") + + +task_map = {} +task_result_map = dict() + +config_file = "config.properties" +# Function to read the variable from the file +def read_global_var(): + if os.path.exists(config_file): + with open(config_file, "r") as f: + for line in f: + if line.startswith("task_type"): + return str(line.split("=")[1].strip()) # Default value is 10 + return 10 + +# Function to write the variable to the file +def write_global_var(value): + with open(config_file, "w") as f: + f.write(f"task_type={value}\n") + + + +# Armorcode Server Routes +@ac_app.before_request +def initialize_tasks(): + global task_map, task_result_map + task_map = TestCases.retrun_task_map() + task_keys = TestCases.return_key_list() + for key in task_keys: + task_map[key] = TaskStatusEnum.IN_QUEUE + +@ac_app.route('/api/http-teleport/get-task', methods=['GET']) +def get_task(): + # Get taskType from query params, default to 'default' if not provided + task_type = read_global_var() + task = task_map.get(task_type, None) + + if task is None: + return jsonify({"error": "No task found"}), 404 + + task_map[task_type] = TaskStatusEnum.IN_PROGRESS + + return jsonify({"data": { + "taskId": task["task_id"], + "url": f"http://localhost:5001{task['url_path']}", + "method": "GET" + }}) + +@ac_app.route('/api/http-teleport/put-result', methods=['POST']) +def put_result(): + print(f"Received result: {request.json}") + + result = request.json + if not result or 'taskId' not in result: + return jsonify({"status": "error", "message": "Invalid result format"}), 400 + + task_id = result['taskId'] + task = next((t for t in TestCases.test_cases if t['task_id'] == task_id), None) + + if not task: + return jsonify({"status": "error", "message": "Task not found"}), 404 + + verification_result, message = verify_task_result(task, result) + + if verification_result == VerificationResult.SUCCESS: + task_map[task_id] = TaskStatusEnum.PASSED + return jsonify({"status": "success", "message": message}), 200 + + else: + task_map[task_id] = TaskStatusEnum.FAILED + return jsonify({"status": "error", "message": message}), 400 + +@ac_app.route('/api/http-teleport/upload-result', methods=['POST']) +def upload_result(): + if 'file' not in request.files: + return jsonify({"error": "No file part"}), 400 + file = request.files['file'] + if file.filename == '': + return jsonify({"error": "No selected file"}), 400 + + if 'task' not in request.form: + return jsonify({"error": "No task data"}), 400 + + task_data = request.form['task'] + + if file: + filename = secure_filename(file.filename) + file.save(os.path.join(ac_app.config['UPLOAD_FOLDER'], filename)) + + # Process the task data as needed + # For example, you might want to save it to a database or perform some operation + + return jsonify({ + "data": f"File {filename} uploaded successfully for task", + "success": True, + "message": "Upload completed" + }), 200 + +@ac_app.route('/api/http-teleport/upload-url', methods=['GET']) +def get_s3_path(): + task_id = request.args.get('fileName', '') + return jsonify({ + "data": { + "putUrl": f"http://localhost:5000/mock-s3-upload?taskId={task_id}", + "getUrl": f"http://localhost:5000/mock-s3-download?taskId={task_id}" + } + }) + +@ac_app.route('/mock-s3-upload', methods=['PUT']) +def mock_s3_upload(): + task_id = request.args.get('taskId') + if not task_id: + return jsonify({"error": "Missing taskId parameter"}), 400 + if 'file' not in request.files: + return jsonify({"error": "No file part"}), 400 + + file = request.files['file'] + if file.filename == '': + return jsonify({"error": "No selected file"}), 400 + + if file: + filename = secure_filename(file.filename) + upload_folder = os.path.join('Resources', 'Files', 'S3_upload') + + if not os.path.exists(upload_folder): + os.makedirs(upload_folder) + + file_path = os.path.join(upload_folder, filename) + file.save(file_path) + + return jsonify({"message": "File uploaded successfully", "filename": filename}), 200 + + + +if __name__ == "__main__": + run_simple('localhost', 5000, ac_app, use_reloader=True, use_debugger=True, use_evalex=True) \ No newline at end of file diff --git a/web-agent/Test/Resources/Files/json_response.json b/web-agent/Test/Resources/Files/json_response.json new file mode 100644 index 0000000..e69de29 diff --git a/web-agent/Test/Resources/TestCases.py b/web-agent/Test/Resources/TestCases.py new file mode 100644 index 0000000..5a38e9f --- /dev/null +++ b/web-agent/Test/Resources/TestCases.py @@ -0,0 +1,59 @@ + +def retrun_task_map(): + return task_map + +def return_key_list(): + return list(task_map.keys()) + +test_cases = [ + { + "task_id": "json_task", + "response_file_path": "json_response.json", + "headers": {"Content-Type": "application/json"}, + "url_path": "/json", + "is_zip": False + }, + { + "task_id": "xml_task", + "response_file_path": "xml_response.xml", + "headers": {"Content-Type": "application/xml"}, + "url_path": "/xml", + "is_zip": False + }, + { + "task_id": "gzip_task", + "response_file_path": "gzip_response.gz", + "headers": {"Content-Type": "application/gzip", "Content-Encoding": "gzip"}, + "url_path": "/gzip", + "is_zip": True + }, + { + "task_id": "big_file_task", + "response_file_path": "big_file.zip", + "headers": {"Content-Type": "application/zip"}, + "url_path": "/big_file", + "is_zip": True + }, + { + "task_id": "timeout_task", + "response_file_path": None, + "headers": {}, + "url_path": "/timeout", + "is_zip": False + }, + { + "task_id": "error_400_task", + "response_file_path": None, + "headers": {}, + "url_path": "/error_400", + "is_zip": False + }, + { + "task_id": "error_500_task", + "response_file_path": None, + "headers": {}, + "url_path": "/error_500", + "is_zip": False + } +] +task_map = {task['task_id']: task for task in test_cases} \ No newline at end of file diff --git a/web-agent/Test/TaskStatus.py b/web-agent/Test/TaskStatus.py new file mode 100644 index 0000000..8c47032 --- /dev/null +++ b/web-agent/Test/TaskStatus.py @@ -0,0 +1,7 @@ +from enum import Enum + +class TaskStatusEnum(Enum): + IN_QUEUE = "In Queue" + IN_PROGRESS = "In Progress" + FAILED = "Failed" + PASSED = "Passed" \ No newline at end of file diff --git a/web-agent/Test/VerifyResult.py b/web-agent/Test/VerifyResult.py new file mode 100644 index 0000000..122919a --- /dev/null +++ b/web-agent/Test/VerifyResult.py @@ -0,0 +1,48 @@ +from enum import Enum + +class VerificationResult(Enum): + SUCCESS = "success" + FAILURE = "failure" + +def verify_task_result(task, result): + verification_functions = { + 'json_task': verify_json_task, + 'xml_task': verify_xml_task, + 'gzip_task': verify_gzip_task, + 'big_file_task': verify_big_file_task, + 'timeout_task': verify_error_task, + 'error_400_task': verify_error_task, + 'error_500_task': verify_error_task + } + + verify_func = verification_functions.get(task['task_id']) + if verify_func: + return verify_func(task, result) + else: + return VerificationResult.FAILURE, "Unknown task type" + +def verify_json_task(task, result): + if 'Content-Type' in result.get('responseHeaders', {}) and result['responseHeaders']['Content-Type'] == 'application/json': + return VerificationResult.SUCCESS, "JSON task verified successfully" + return VerificationResult.FAILURE, "JSON task verification failed" + +def verify_xml_task(task, result): + if 'Content-Type' in result.get('responseHeaders', {}) and result['responseHeaders']['Content-Type'] == 'application/xml': + return VerificationResult.SUCCESS, "XML task verified successfully" + return VerificationResult.FAILURE, "XML task verification failed" + +def verify_gzip_task(task, result): + if 'Content-Encoding' in result.get('responseHeaders', {}) and result['responseHeaders']['Content-Encoding'] == 'gzip': + return VerificationResult.SUCCESS, "GZIP task verified successfully" + return VerificationResult.FAILURE, "GZIP task verification failed" + +def verify_big_file_task(task, result): + # Add specific verification for big file task if needed + return VerificationResult.SUCCESS, "Big file task received" + +def verify_error_task(task, result): + expected_status = int(task['task_id'].split('_')[1]) + if 'status' in result and result['status'] == expected_status: + return VerificationResult.SUCCESS, f"{task['task_id']} verified successfully" + return VerificationResult.FAILURE, f"{task['task_id']} verification failed" + diff --git a/web-agent/Test/config.properties b/web-agent/Test/config.properties new file mode 100644 index 0000000..33a6d8c --- /dev/null +++ b/web-agent/Test/config.properties @@ -0,0 +1 @@ +task_type=json_task diff --git a/web-agent/app/worker.py b/web-agent/app/worker.py index ec92a0c..05a0d3f 100644 --- a/web-agent/app/worker.py +++ b/web-agent/app/worker.py @@ -268,8 +268,8 @@ def process_task(task: Dict[str, Any]) -> Dict[str, Any]: try: # Running the request - # timeout = round((expiryTime - round(time.time() * 1000)) / 1000) - # logger.info("expiry %s, %s", expiryTime, timeout) + timeout = round((expiryTime - round(time.time() * 1000)) / 1000) + logger.info("expiry %s, %s", expiryTime, timeout) logger.debug("Request for task %s with headers %s and input_data %s", taskId, headers, input_data) check_and_update_encode_url(headers, url) From 8ac017ddf47047648b1bb79d46de73f750527ee0 Mon Sep 17 00:00:00 2001 From: dmeenaarmorcode Date: Wed, 5 Feb 2025 21:34:15 +0530 Subject: [PATCH 2/3] Long_term fix --- web-agent/app/worker.py | 22 +++------------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/web-agent/app/worker.py b/web-agent/app/worker.py index 05a0d3f..2b7d087 100644 --- a/web-agent/app/worker.py +++ b/web-agent/app/worker.py @@ -57,8 +57,7 @@ def main() -> None: parser.add_argument("--serverUrl", required=False, help="Server Url") parser.add_argument("--apiKey", required=False, help="Api Key") parser.add_argument("--index", required=False, help="Agent index no", default="_prod") - parser.add_argument("--timeout", required=False, help="timeout", default=30) - parser.add_argument("--verify", required=False, help="Verify Cert", default=True) + parser.add_argument("--verify", action="store_true", help="Verify Cert", default=False) parser.add_argument("--debugMode", required=False, help="Enable debug Mode", default=True) parser.add_argument("--inwardProxyHttps", required=False, help="Pass inward Https proxy", default=None) @@ -74,8 +73,7 @@ def main() -> None: server_url = args.serverUrl api_key = args.apiKey agent_index: str = args.index - timeout_cmd = args.timeout - verify_cmd = args.verify + verify_cert = args.verify debug_cmd = args.debugMode upload_to_ac = args.uploadToAc @@ -108,13 +106,6 @@ def main() -> None: if str(debug_cmd).lower() == "false": debug_mode = False - if verify_cmd is not None: - if str(verify_cmd).lower() == "false": - verify_cert = False - - if timeout_cmd is not None: - timeout = int(timeout_cmd) - if os.getenv('verify') is not None: if str(os.getenv('verify')).lower() == "false": verify_cert = False @@ -137,14 +128,7 @@ def main() -> None: logger.error("Empty serverUrl %s", server_url) raise ValueError("Server URL and API Key must be provided either as arguments or environment variables") - # Creating thread pool to use other thread if one thread is blocked in I/O - # pool: concurrent.futures.ThreadPoolExecutor = concurrent.futures.ThreadPoolExecutor(max_workers=2) - # pool.submit(process) - # pool.submit(process) - # - # pool.shutdown(wait=True) - # Instantiate RateLimiter for 25 requests per 15 seconds window rate_limiter = RateLimiter(request_limit=25, time_window=15) process() @@ -274,7 +258,7 @@ def process_task(task: Dict[str, Any]) -> Dict[str, Any]: logger.debug("Request for task %s with headers %s and input_data %s", taskId, headers, input_data) check_and_update_encode_url(headers, url) response: requests.Response = requests.request(method, url, headers=headers, data=input_data, stream=True, - timeout=timeout, verify=verify_cert, proxies=inward_proxy) + timeout=(7, timeout), verify=verify_cert, proxies=inward_proxy) logger.info("Response: %d", response.status_code) data: Any = None From d0882e8356fbe02f5bedc7ade74e172cbf484444 Mon Sep 17 00:00:00 2001 From: dmeenaarmorcode Date: Wed, 5 Feb 2025 21:35:07 +0530 Subject: [PATCH 3/3] test_commit --- web-agent/README.md | 6 +-- web-agent/Test/AgentIntegrationTest.py | 28 ++++++------- web-agent/Test/Mock_ArmorCode_server.py | 16 +------- .../Test/{VerifyResult.py => TestHelper.py} | 41 +++++++++++++++++++ web-agent/Test/config.properties | 2 +- 5 files changed, 61 insertions(+), 32 deletions(-) rename web-agent/Test/{VerifyResult.py => TestHelper.py} (63%) diff --git a/web-agent/README.md b/web-agent/README.md index 52ca8ec..3fe441e 100644 --- a/web-agent/README.md +++ b/web-agent/README.md @@ -96,11 +96,11 @@ Steps for customer 3. Run command: ```commandline - python3 worker.py --serverUrl 'https://app.armorcode.com' --apiKey `` --index 0 --timeout 30 + python3 worker.py --serverUrl 'https://app.armorcode.com' --apiKey `` ``` -4. If you don't want to do certificates validations (needed in case if VM don't have any certificates assigned and making https request) pass this extra argument at the end +4. If you want to do certificates validations (needed in case if API end point requires valid certificates) pass this extra argument at the end ```commandline - --verify=False + --verify ``` 5. If you have HTTPS proxy to make calls to ArmorCode API, add this argument. ex ## diff --git a/web-agent/Test/AgentIntegrationTest.py b/web-agent/Test/AgentIntegrationTest.py index da19df9..c2d1d43 100644 --- a/web-agent/Test/AgentIntegrationTest.py +++ b/web-agent/Test/AgentIntegrationTest.py @@ -10,6 +10,7 @@ import Mock_ArmorCode_server from MockClientServer import client_app from Resources import TestCases +from TestHelper import state_manager file_directory = "Resources/Files/FilesToUpload" @@ -37,18 +38,18 @@ def setUpClass(cls): # Install requirements subprocess.run(['pip3', 'install', '-r', 'requirements.txt']) - # Start Mock_ArmorCode_server - # cls.ac_thread = threading.Thread(target=lambda: ac_app.run(port=5000)) - # cls.ac_thread.daemon = True - # cls.ac_thread.start() - # - # # Start MockClientServer - # cls.client_thread = threading.Thread(target=lambda: client_app.run(port=5001)) - # cls.client_thread.daemon = True - # cls.client_thread.start() - # - # # Wait for servers to start - # time.sleep(2) + #Start Mock_ArmorCode_server + cls.ac_thread = threading.Thread(target=lambda: Mock_ArmorCode_server.ac_app.run(port=5000)) + cls.ac_thread.daemon = True + cls.ac_thread.start() + + # Start MockClientServer + cls.client_thread = threading.Thread(target=lambda: client_app.run(port=5001)) + cls.client_thread.daemon = True + cls.client_thread.start() + + # Wait for servers to start + time.sleep(2) @staticmethod def create_large_zip_file(file_path, size_bytes): @@ -87,8 +88,7 @@ def test_tasks(self): for case in TestCases.test_cases: with self.subTest(task_id=case['task_id']): # Set task type to return - Mock_ArmorCode_server.write_global_var(case['task_id']) - + state_manager.set_task_type(case['task_id']) # Run worker and wait for result result = self.run_worker_and_wait(case['task_id']) diff --git a/web-agent/Test/Mock_ArmorCode_server.py b/web-agent/Test/Mock_ArmorCode_server.py index 53fdfdd..b45ee42 100644 --- a/web-agent/Test/Mock_ArmorCode_server.py +++ b/web-agent/Test/Mock_ArmorCode_server.py @@ -9,7 +9,7 @@ from werkzeug.utils import secure_filename from TaskStatus import TaskStatusEnum -from VerifyResult import VerificationResult, verify_task_result +from TestHelper import VerificationResult, verify_task_result, state_manager import HttpTeleportTask @@ -22,18 +22,6 @@ config_file = "config.properties" # Function to read the variable from the file -def read_global_var(): - if os.path.exists(config_file): - with open(config_file, "r") as f: - for line in f: - if line.startswith("task_type"): - return str(line.split("=")[1].strip()) # Default value is 10 - return 10 - -# Function to write the variable to the file -def write_global_var(value): - with open(config_file, "w") as f: - f.write(f"task_type={value}\n") @@ -49,7 +37,7 @@ def initialize_tasks(): @ac_app.route('/api/http-teleport/get-task', methods=['GET']) def get_task(): # Get taskType from query params, default to 'default' if not provided - task_type = read_global_var() + task_type = state_manager.get_task_type() task = task_map.get(task_type, None) if task is None: diff --git a/web-agent/Test/VerifyResult.py b/web-agent/Test/TestHelper.py similarity index 63% rename from web-agent/Test/VerifyResult.py rename to web-agent/Test/TestHelper.py index 122919a..4efc021 100644 --- a/web-agent/Test/VerifyResult.py +++ b/web-agent/Test/TestHelper.py @@ -1,5 +1,46 @@ from enum import Enum +import json +import os + +import json +import os + +class TestStateManager: + _file_path = 'Resources/Files/test_state.json' + + def get_task_type(self): + if not os.path.exists(self._file_path): + return "default_value" + with open(self._file_path, 'r') as f: + data = json.load(f) + return data.get("task_type", "default_value") + + def set_task_type(self, task_type): + data = {} + # Load existing data if the file exists + if os.path.exists(self._file_path): + with open(self._file_path, 'r') as f: + try: + data = json.load(f) + except json.JSONDecodeError: + # Handle empty or malformed JSON file + data = {} + # Update the task_type key + data["task_type"] = task_type + # Write back the updated data + with open(self._file_path, 'w') as f: + json.dump(data, f, indent=4) + + def reset_state(self): + if os.path.exists(self._file_path): + os.remove(self._file_path) + # Create an empty JSON file + with open(self._file_path, 'w') as f: + json.dump({}, f, indent=4) + +state_manager = TestStateManager() + class VerificationResult(Enum): SUCCESS = "success" FAILURE = "failure" diff --git a/web-agent/Test/config.properties b/web-agent/Test/config.properties index 33a6d8c..9619524 100644 --- a/web-agent/Test/config.properties +++ b/web-agent/Test/config.properties @@ -1 +1 @@ -task_type=json_task +task_type=error_500_task