From b09f4de5f30541689aa8490cc1f6a0a936a38ac9 Mon Sep 17 00:00:00 2001 From: Praveen Gupta Date: Thu, 14 Nov 2024 00:39:38 -0800 Subject: [PATCH] Adds test cases --- .gitignore | 10 ++-- README.md | 6 +-- tests/README.md | 2 + tests/__init__.py | 0 tests/sample_app/growlithe_config.yaml | 6 +++ tests/sample_app/src/function1.py | 26 ++++++++++ tests/sample_app/src/function2.py | 26 ++++++++++ tests/sample_app/state_machine.asl.json | 53 ++++++++++++++++++++ tests/sample_app/template.yaml | 40 ++++++++++++++++ tests/test_cli.py | 63 ++++++++++++++++++++++++ tests/test_graph.py | 34 +++++++++++++ tests/test_policy.py | 64 +++++++++++++++++++++++++ 12 files changed, 324 insertions(+), 6 deletions(-) create mode 100644 tests/README.md create mode 100644 tests/__init__.py create mode 100644 tests/sample_app/growlithe_config.yaml create mode 100644 tests/sample_app/src/function1.py create mode 100644 tests/sample_app/src/function2.py create mode 100644 tests/sample_app/state_machine.asl.json create mode 100644 tests/sample_app/template.yaml create mode 100644 tests/test_cli.py create mode 100644 tests/test_graph.py create mode 100644 tests/test_policy.py diff --git a/.gitignore b/.gitignore index d0c6210..5a64009 100644 --- a/.gitignore +++ b/.gitignore @@ -5,7 +5,8 @@ # codeql databases codeqldb/ - +growlithe_*/ +benchmarks/*/growlithe*.yaml # test files test.ql test.py @@ -37,5 +38,8 @@ __pycache__/ tmp/ policies.txt -env/ -growlithe.egg-info/ \ No newline at end of file +venv/ +growlithe.egg-info/ + +# Growlithe outputs +growlithe_SampleApp/ \ No newline at end of file diff --git a/README.md b/README.md index fdf5365..6634536 100644 --- a/README.md +++ b/README.md @@ -8,8 +8,8 @@ Our 2025 IEEE S&P paper provides more details about the design of Growlithe: ## Setup - Create a [new virtual environment](https://docs.python.org/3/library/venv.html) with python v3.10, and activate it. ```bash -python -m venv env -source env/bin/activate # On Windows use `env\Scripts\activate` +python -m venv venv +source venv/bin/activate # On Windows use `venv\Scripts\activate` ``` - Install CodeQL and dependencies by following [/growlithe/graph/codeql/README.md](/growlithe/graph/codeql/README.md) @@ -24,7 +24,7 @@ source env/bin/activate # On Windows use `env\Scripts\activate` Activate the virtual environment, then: - Navigate to your serverless application, create a file `growlithe_config.yaml` with the following configuration: ```yaml -app_path: +# app_path: app_name: src_dir: app_config_path: diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..49bbbab --- /dev/null +++ b/tests/README.md @@ -0,0 +1,2 @@ +Run all unit tests by using: +`python -m unittest -v` \ No newline at end of file diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/sample_app/growlithe_config.yaml b/tests/sample_app/growlithe_config.yaml new file mode 100644 index 0000000..e4a0174 --- /dev/null +++ b/tests/sample_app/growlithe_config.yaml @@ -0,0 +1,6 @@ +app_path: . +app_name: SampleApp +src_dir: src +app_config_path: template.yaml +app_config_type: SAM +cloud_provider: AWS \ No newline at end of file diff --git a/tests/sample_app/src/function1.py b/tests/sample_app/src/function1.py new file mode 100644 index 0000000..7efc217 --- /dev/null +++ b/tests/sample_app/src/function1.py @@ -0,0 +1,26 @@ +import json +import time +import os +import boto3 + + +def lambda_handler(event, context): + body = event["body"] + + s3 = boto3.resource("s3") + bucket_name = "sample-test-bucket" + bucket = s3.Bucket(bucket_name) + object_key = os.getenv("AWS_LAMBDA_FUNCTION_NAME") + + tempFile = "/tmp/" + object_key + os.makedirs(os.path.dirname(tempFile), exist_ok=True) + + bucket.download_file(object_key, tempFile) + + time.sleep(0.065) + + output_key = f"output/{object_key}" + bucket.upload_file(tempFile, output_key) + + response = {"statusCode": 200, "body": body} + return response diff --git a/tests/sample_app/src/function2.py b/tests/sample_app/src/function2.py new file mode 100644 index 0000000..7efc217 --- /dev/null +++ b/tests/sample_app/src/function2.py @@ -0,0 +1,26 @@ +import json +import time +import os +import boto3 + + +def lambda_handler(event, context): + body = event["body"] + + s3 = boto3.resource("s3") + bucket_name = "sample-test-bucket" + bucket = s3.Bucket(bucket_name) + object_key = os.getenv("AWS_LAMBDA_FUNCTION_NAME") + + tempFile = "/tmp/" + object_key + os.makedirs(os.path.dirname(tempFile), exist_ok=True) + + bucket.download_file(object_key, tempFile) + + time.sleep(0.065) + + output_key = f"output/{object_key}" + bucket.upload_file(tempFile, output_key) + + response = {"statusCode": 200, "body": body} + return response diff --git a/tests/sample_app/state_machine.asl.json b/tests/sample_app/state_machine.asl.json new file mode 100644 index 0000000..de21874 --- /dev/null +++ b/tests/sample_app/state_machine.asl.json @@ -0,0 +1,53 @@ +{ + "Comment": "A linear chain of Lambda functions", + "StartAt": "function1", + "States": { + "function1": { + "Type": "Task", + "Resource": "arn:aws:states:::lambda:invoke", + "OutputPath": "$.Payload", + "Parameters": { + "Payload.$": "$", + "FunctionName": "${function1-arn}" + }, + "Retry": [ + { + "ErrorEquals": [ + "Lambda.ServiceException", + "Lambda.AWSLambdaException", + "Lambda.SdkClientException", + "Lambda.TooManyRequestsException" + ], + "IntervalSeconds": 1, + "MaxAttempts": 3, + "BackoffRate": 2 + } + ], + "Next": "function2", + "End": false + }, + "function2": { + "Type": "Task", + "Resource": "arn:aws:states:::lambda:invoke", + "OutputPath": "$.Payload", + "Parameters": { + "Payload.$": "$", + "FunctionName": "${function2-arn}" + }, + "Retry": [ + { + "ErrorEquals": [ + "Lambda.ServiceException", + "Lambda.AWSLambdaException", + "Lambda.SdkClientException", + "Lambda.TooManyRequestsException" + ], + "IntervalSeconds": 1, + "MaxAttempts": 3, + "BackoffRate": 2 + } + ], + "End": true + } + } +} \ No newline at end of file diff --git a/tests/sample_app/template.yaml b/tests/sample_app/template.yaml new file mode 100644 index 0000000..bc0a58d --- /dev/null +++ b/tests/sample_app/template.yaml @@ -0,0 +1,40 @@ +AWSTemplateFormatVersion: '2010-09-09' +Description: Serverless State Machine App +Resources: + Function1: + Properties: + CodeUri: src/ + Environment: + Variables: + BUCKET_NAME: sample-test-bucket + Handler: function1.lambda_handler + Policies: + - S3CrudPolicy: + BucketName: sample-test-bucket + Runtime: python3.10 + Type: AWS::Serverless::Function + Function2: + Properties: + CodeUri: src/ + Environment: + Variables: + BUCKET_NAME: sample-test-bucket + Handler: function2.lambda_handler + Policies: + - S3CrudPolicy: + BucketName: sample-test-bucket + Runtime: python3.10 + Type: AWS::Serverless::Function + StateMachine: + Properties: + DefinitionSubstitutions: + function1-arn: !GetAtt 'Function1.Arn' + function2-arn: !GetAtt 'Function2.Arn' + DefinitionUri: state_machine.asl.json + Policies: + - LambdaInvokePolicy: + FunctionName: '*' + - S3CrudPolicy: + BucketName: sample-test-bucket + Type: AWS::Serverless::StateMachine +Transform: AWS::Serverless-2016-10-31 \ No newline at end of file diff --git a/tests/test_cli.py b/tests/test_cli.py new file mode 100644 index 0000000..7ad9f05 --- /dev/null +++ b/tests/test_cli.py @@ -0,0 +1,63 @@ +import os +import unittest +from click.testing import CliRunner +from growlithe.cli.cli import cli + +DEBUG = True + + +def print_runner_logs(result): + if DEBUG: + print(f"Exit code: {result.exit_code}") + print(f"Output: {result.output}") + print(f"Exception: {result.exception}") + + +class TestCli(unittest.TestCase): + @classmethod + def setUpClass(self): + self.runner = CliRunner() + self.original_dir = os.getcwd() + + # Get the directory of the current file + current_dir = os.path.dirname(os.path.abspath(__file__)) + # Construct the path to the sample_app directory + self.sample_app_dir = os.path.join(current_dir, "sample_app") + # Construct the path to growlithe_config.yaml + self.custom_config_path = os.path.join( + self.sample_app_dir, "growlithe_config.yaml" + ) + + print(f"Current working directory: {os.getcwd()}") + print(f"Using config path: {self.custom_config_path}") + print(f"Config file exists: {os.path.exists(self.custom_config_path)}") + print( + "----------------------------------------------------------------------\n" + ) + # Change the current working directory to sample_app + os.chdir(self.sample_app_dir) + + def tearDown(self): + os.chdir(self.original_dir) + + def test_analyze_with_custom_config(self): + # Run the CLI command with the custom config option + result = self.runner.invoke( + cli, ["--config", self.custom_config_path, "analyze"] + ) + print_runner_logs(result) + + # Check that the command executed successfully + self.assertEqual(result.exit_code, 0) + + def test_apply_with_custom_config(self): + # Run the CLI command with the custom config option + result = self.runner.invoke(cli, ["--config", self.custom_config_path, "apply"]) + print_runner_logs(result) + + # Check that the command executed successfully + self.assertEqual(result.exit_code, 0) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_graph.py b/tests/test_graph.py new file mode 100644 index 0000000..1cd78c1 --- /dev/null +++ b/tests/test_graph.py @@ -0,0 +1,34 @@ +import os +import unittest +from growlithe.cli.analyze import analyze +from growlithe.config import get_config + + +class TestGraph(unittest.TestCase): + @classmethod + def setUpClass(self): + # Get the directory of the current file + current_dir = os.path.dirname(os.path.abspath(__file__)) + # Construct the path to the sample_app directory + self.sample_app_dir = os.path.join(current_dir, "sample_app") + # Construct the path to growlithe_config.yaml + self.custom_config_path = os.path.join( + self.sample_app_dir, "growlithe_config.yaml" + ) + + self.config = get_config(os.path.abspath(self.custom_config_path)) + + def tearDown(self): + pass + + def test_graph(self): + graph = analyze(self.config) + print(len(graph.nodes)) + print(len(graph.functions)) + + self.assertEqual(len(graph.functions), 2) + self.assertEqual(len(graph.edges), 7) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_policy.py b/tests/test_policy.py new file mode 100644 index 0000000..83c537e --- /dev/null +++ b/tests/test_policy.py @@ -0,0 +1,64 @@ +import os, re, json +import unittest +from growlithe.cli.apply import apply +from growlithe.config import get_config + + +def match_pattern(value, pattern): + # Create a regex pattern that matches any number followed by the rest of the string + return value.partition(":")[2].strip() == pattern + # return re.match(regex_pattern, value) is not None + + +def add_policy(policy_file_path, policy_str): + # Read the JSON file + with open(policy_file_path, "r") as file: + data = json.load(file) + + # Find the specific entry and update it + for entry in data: + if match_pattern(entry["source"], "tempfs:$tempFile") and match_pattern( + entry["sink"], "sample-test-bucket:$output_key" + ): + # Update the read policy + entry["write"] = policy_str + print(f"Updated policy in : {entry}") + break + else: + print("Edge not found.") + + # Write the updated data back to the JSON file + with open(policy_file_path, "w") as file: + json.dump(data, file, indent=4) + + +class TestPolicy(unittest.TestCase): + @classmethod + def setUpClass(self): + # Get the directory of the current file + current_dir = os.path.dirname(os.path.abspath(__file__)) + + # Construct the path to the sample_app directory + self.sample_app_dir = os.path.join(current_dir, "sample_app") + + # Construct the path to growlithe_config.yaml + self.custom_config_path = os.path.join( + self.sample_app_dir, "growlithe_config.yaml" + ) + + self.config = get_config(os.path.abspath(self.custom_config_path)) + self.policy_file_path = os.path.join( + self.sample_app_dir, "growlithe_SampleApp", "policy_spec.json" + ) + + def tearDown(self): + pass + + def test_policy(self): + policy = "eq(ResourceRegion, 'us-west-1')" + add_policy(self.policy_file_path, policy) + apply(self.config) + + +if __name__ == "__main__": + unittest.main(verbosity=2)