From ffd4681cbd51c0c8bdd48a2dc2b7879db0f849f3 Mon Sep 17 00:00:00 2001
From: Hope Aitchison <hope.aitchison@digital.justice.gov.uk>
Date: Mon, 18 Dec 2023 16:02:49 +0000
Subject: [PATCH] updating unit test for subnet and sg ids and adding in zip
 package type option

---
 main.tf                   |  2 ++
 test/lambda_test.go       |  6 ++++++
 test/unit-test/main.tf    | 39 ++++++++++++++++++++++++++++++++++++++-
 test/unit-test/outputs.tf |  8 ++++++++
 variables.tf              | 10 ++++++++--
 5 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/main.tf b/main.tf
index 27e15b3..1765046 100644
--- a/main.tf
+++ b/main.tf
@@ -65,10 +65,12 @@ resource "aws_lambda_function" "this" { #tfsec:ignore:aws-lambda-enable-tracing
   description                    = var.description
   reserved_concurrent_executions = var.reserved_concurrent_executions
   image_uri                      = var.image_uri
+  filename                       = var.filename
   package_type                   = var.package_type
   role                           = var.create_role ? aws_iam_role.this[0].arn : var.lambda_role
   timeout                        = var.timeout
   memory_size                    = var.memory_size
+  
   dynamic "tracing_config" {
     for_each = var.tracing_mode != null ? [1] : []
     content {
diff --git a/test/lambda_test.go b/test/lambda_test.go
index 7c70354..b4568bb 100644
--- a/test/lambda_test.go
+++ b/test/lambda_test.go
@@ -20,7 +20,13 @@ func TestLambdaCreation(t *testing.T) {
 
 	functionName := terraform.Output(t, terraformOptions, "function_name")
 	resultCode := terraform.Output(t, terraformOptions, "result_code")
+	checkSubnetID := terraform.Output(t, terraformOptions, "subnet_id")
+	checkSecurityGroupId := terraform.Output(t, terraformOptions, "security_group_id")
 
 	assert.Regexp(t, regexp.MustCompile(`^instance-scheduler-lambda-function*`), functionName)
 	assert.Regexp(t, regexp.MustCompile(`^200*`), resultCode)
+	assert.Regexp(t, regexp.MustCompile(`^subnet-*`), checkSubnetId)
+	assert.Regexp(t, regexp.MustCompile(`^sg-*`), checkSecurityGroupId)
 }
+
+func
diff --git a/test/unit-test/main.tf b/test/unit-test/main.tf
index b2ebf11..fd8de13 100644
--- a/test/unit-test/main.tf
+++ b/test/unit-test/main.tf
@@ -30,6 +30,20 @@ module "module_test" {
 
 }
 
+module "module_vpc_test" {
+  source                  = "../../"
+  application_name        = local.application_name
+  description             = "vpc attached test lambda"
+  tags                    = local.tags
+  function_name           = "vpc-attached-lambda-function"
+  vpc_subnet_ids          = aws_subnet.lambda_subnet_test.id
+  vpc_security_group_ids  = aws_security_group.lambda_security_group_test.id
+  create_role             = true
+  role_name               = "InstanceSchedulerLambdaFunctionPolicyVPCTest"
+  policy_json_attached    = true
+  policy_json             = data.aws_iam_policy_document.instance-scheduler-lambda-function-policy.json
+}
+
 resource "aws_cloudwatch_event_rule" "instance_scheduler_weekly_stop_at_night" {
   name                = "instance_scheduler_weekly_stop_at_night"
   description         = "Call Instance Scheduler with Stop action at 8:00 pm (UTC) every Monday through Friday"
@@ -156,4 +170,27 @@ resource "aws_lambda_invocation" "test_invocation" {
     {
       action = "Test"
   })
-}
\ No newline at end of file
+}
+
+resource "aws_vpc" "lambda_vpc_config_test"{
+  cidr_block = "10.0.0.0/16"
+}
+
+resource "aws_subnet" "lambda_subnet_test"{
+  vpc_id     = aws_vpc.lambda_vpc_config_test.id
+  cidr_block = "10.0.1.0/24"
+}
+
+resource "aws_security_group" "lambda_security_group_test" {
+  name        = "lambda-vpc-test"
+  description = "lambda attached to vpc test security group"
+  vpc_id = aws_vpc.lambda_vpc_config_test.id
+
+  egress {
+    from_port        = 0
+    to_port          = 0
+    protocol         = "-1"
+  }
+
+  tags = local.tags
+}
diff --git a/test/unit-test/outputs.tf b/test/unit-test/outputs.tf
index c387d4f..1f5bb9b 100644
--- a/test/unit-test/outputs.tf
+++ b/test/unit-test/outputs.tf
@@ -8,4 +8,12 @@ output "function_name" {
 
 output "result_code" {
   value = jsondecode(aws_lambda_invocation.test_invocation.result)["statusCode"]
+}
+
+output "security_group_id" {
+  value = module.module_test.lambda_function_security_group_id
+}
+
+output "subnet_id" {
+  value = module.module_test.lambda_function_subnet_id
 }
\ No newline at end of file
diff --git a/variables.tf b/variables.tf
index 6837a21..54dae7d 100644
--- a/variables.tf
+++ b/variables.tf
@@ -76,9 +76,9 @@ variable "environment_variables" {
 }
 
 variable "package_type" {
-  description = "The Lambda deployment package type. Valid options: Image"
+  description = "The Lambda deployment package type. Valid options: Image or Zip"
   type        = string
-  default     = "Image"
+  default     = "Zip"
 }
 
 variable "image_uri" {
@@ -87,6 +87,12 @@ variable "image_uri" {
   default     = null
 }
 
+variable "filename" {
+  description = "The absolute path to an existing zip-file to use"
+  type        = string
+  default     = null
+}
+
 variable "timeout" {
   description = "The amount of time your Lambda Function has to run in seconds."
   type        = number