diff --git a/test/unit-tests/aws/api-gateway-v1-wrapper.test.ts b/test/unit-tests/aws/api-gateway-v1-wrapper.test.ts
index e6e3093e..1d9b2f16 100644
--- a/test/unit-tests/aws/api-gateway-v1-wrapper.test.ts
+++ b/test/unit-tests/aws/api-gateway-v1-wrapper.test.ts
@@ -339,6 +339,14 @@ describe("API Gateway V1 wrapper checks", () => {
           restApiId: "test_rest_api_id",
           basePath: "test",
           stage: "test"
+        }, {
+          restApiId: "test_rest_api_id2",
+          basePath: "test2",
+          stage: "test"
+        }, {
+          restApiId: "test_rest_api_id3",
+          basePath: "test3",
+          stage: "dummy"
         }]
       });
 
@@ -348,8 +356,10 @@ describe("API Gateway V1 wrapper checks", () => {
       }));
 
       const actualResult = await apiGatewayV1Wrapper.getBasePathMappings(dc);
+      // should be filtered by stage
       const expectedResult = [
-        new ApiGatewayMap("test_rest_api_id", "test", "test", null)
+        new ApiGatewayMap("test_rest_api_id", "test", "test", null),
+        new ApiGatewayMap("test_rest_api_id2", "test2", "test", null)
       ];
 
       expect(actualResult).to.eql(expectedResult);
diff --git a/test/unit-tests/aws/api-gateway-v2-wrapper.test.ts b/test/unit-tests/aws/api-gateway-v2-wrapper.test.ts
index f2f63bbc..ac2d7efc 100644
--- a/test/unit-tests/aws/api-gateway-v2-wrapper.test.ts
+++ b/test/unit-tests/aws/api-gateway-v2-wrapper.test.ts
@@ -372,6 +372,16 @@ describe("API Gateway V2 wrapper checks", () => {
           ApiMappingKey: "test",
           Stage: "test",
           ApiMappingId: "test_id"
+        }, {
+          ApiId: "test_rest_api_id2",
+          ApiMappingKey: "test2",
+          Stage: "test",
+          ApiMappingId: "test_id2"
+        }, {
+          ApiId: "test_rest_api_id3",
+          ApiMappingKey: "test3",
+          Stage: "dummy",
+          ApiMappingId: "test_id3"
         }]
       });
 
@@ -381,8 +391,10 @@ describe("API Gateway V2 wrapper checks", () => {
       }));
 
       const actualResult = await apiGatewayV2Wrapper.getBasePathMappings(dc);
+      // should be filtered by stage
       const expectedResult = [
-        new ApiGatewayMap("test_rest_api_id", "test", "test", "test_id")
+        new ApiGatewayMap("test_rest_api_id", "test", "test", "test_id"),
+        new ApiGatewayMap("test_rest_api_id2", "test2", "test", "test_id2")
       ];
 
       expect(actualResult).to.eql(expectedResult);