Adds test cases

.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/
+venv/
+growlithe.egg-info/
+
+# Growlithe outputs
+growlithe_SampleApp/

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: <Relative path to the main application>
+# app_path: <Relative path to the main application>
 app_name: <Name of the application>
 src_dir: <Source code of the application relative to app_path>
 app_config_path: <Relative path to the application configuration>

tests/README.md

@@ -0,0 +1,2 @@
+Run all unit tests by using:
+`python -m unittest -v`

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

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

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

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
+    }
+  }
+}

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

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)

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)

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)