diff --git a/main.tf b/main.tf
index 370c071..f8e0bcd 100644
--- a/main.tf
+++ b/main.tf
@@ -85,10 +85,10 @@ resource "aws_lambda_function" "this" { #tfsec:ignore:aws-lambda-enable-tracing
     }
   }
   dynamic "vpc_config" {
-    for_each = var.vpc_subnet_ids != null && var.vpc_security_group_ids != null ? [true] : []
+    for_each = var.subnet_ids != null && var.security_group_ids != null ? [true] : []
     content {
-      security_group_ids = var.vpc_security_group_ids
-      subnet_ids         = var.vpc_subnet_ids
+      security_group_ids = var.security_group_ids
+      subnet_ids         = var.subnet_ids
     }
   }
 }
diff --git a/test/unit-test/main.tf b/test/unit-test/main.tf
index 6bf9f88..ebc8a4a 100644
--- a/test/unit-test/main.tf
+++ b/test/unit-test/main.tf
@@ -39,12 +39,15 @@ module "module_lambda_vpc_test" {
   source_code_hash       = data.archive_file.lambda-zip.output_base64sha256
   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
+
+  vpc_config {
+    subnet_ids           = aws_subnet.lambda_subnet_test.id
+    security_group_ids   = aws_security_group.lambda_security_group_test.id
+  }
 }
 
 resource "aws_cloudwatch_event_rule" "instance_scheduler_weekly_stop_at_night" {
diff --git a/test/unit-test/outputs.tf b/test/unit-test/outputs.tf
index 750ad56..48361ea 100644
--- a/test/unit-test/outputs.tf
+++ b/test/unit-test/outputs.tf
@@ -11,9 +11,9 @@ output "result_code" {
 }
 
 output "security_group_id" {
-  value = module.module_lambda_vpc_test.vpc_security_group_ids
+  value = module.module_lambda_vpc_test.security_group_ids
 }
 
 output "subnet_id" {
-  value = module.module_lambda_vpc_test.vpc_subnet_ids
+  value = module.module_lambda_vpc_test.subnet_ids
 }
\ No newline at end of file
diff --git a/variables.tf b/variables.tf
index 1766e10..ecfc1db 100644
--- a/variables.tf
+++ b/variables.tf
@@ -151,13 +151,13 @@ variable "sns_topic_on_success" {
   default     = ""
 }
 
-variable "vpc_subnet_ids" {
+variable "subnet_ids" {
   description = "List of subnet ids when Lambda Function should run in the VPC. Usually private or intra subnets."
   type        = list(string)
   default     = null
 }
 
-variable "vpc_security_group_ids" {
+variable "security_group_ids" {
   description = "List of security group ids when Lambda Function should run in the VPC."
   type        = list(string)
   default     = null