TDR-3426 - Set limit for numberofrefs (#31)

src/main/resources/application.conf

@@ -9,4 +9,6 @@ dynamodb {
   referenceCounter = ${?REFERENCE_COUNTER}
   queryParam = "numberofrefs"
   queryParam = ${?QUERY_PARAM}
+  referenceLimit = 5000
+  referenceLimit = ${?REFERENCE_LIMIT}
 }

src/main/scala/uk/gov/nationalarchives/referencegenerator/Lambda.scala

@@ -17,15 +17,19 @@ import scala.util.{Failure, Success, Try}
 class Lambda extends RequestHandler[APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent] {
 
   override def handleRequest(event: APIGatewayProxyRequestEvent, context: Context): APIGatewayProxyResponseEvent = {
+    val limit = config.getString("dynamodb.referenceLimit").toInt
     val queryParams = event.getQueryStringParameters
     val queryParam: String = config.getString("dynamodb.queryParam")
-    val convertQueryToInt: Try[Int] = Try(queryParams.get(queryParam).toInt)
-    convertQueryToInt match {
+    val numberOfRefsQuery: Try[Int] = Try {
+      val numberOfRefs = queryParams.get(queryParam).toInt
+      if (numberOfRefs > limit) throw new IllegalArgumentException(s"$queryParam is greater than $limit") else numberOfRefs
+    }
+    numberOfRefsQuery match {
       case Success(numberOfReferences) =>
         process(Input(numberOfReferences))
       case Failure(exception) =>
         val response = new APIGatewayProxyResponseEvent()
-        logger.error(exception.getMessage)
+        logger.error(exception.toString)
         response.setStatusCode(500)
         response.setBody(exception.getMessage)
         response

src/test/resources/application.conf

@@ -4,4 +4,5 @@ dynamodb {
   keyVal = "fileCounter"
   referenceCounter = "referenceCounter"
   queryParam = "numberofrefs"
+  referenceLimit = 5000
 }

src/test/scala/uk/gov/nationalarchives/LambdaTest.scala

@@ -1,18 +1,35 @@
 package uk.gov.nationalarchives
 
-import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyResponseEvent
+import com.amazonaws.services.lambda.runtime.events.{APIGatewayProxyRequestEvent, APIGatewayProxyResponseEvent}
+import com.amazonaws.services.lambda.runtime.{ClientContext, CognitoIdentity, Context, LambdaLogger}
 import com.dimafeng.testcontainers.DynaliteContainer
 import com.typesafe.config.{Config, ConfigFactory, ConfigValueFactory}
 import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
 import uk.gov.nationalarchives.referencegenerator.Lambda
 import uk.gov.nationalarchives.utils.TestContainerUtils
 
+import scala.jdk.CollectionConverters.MapHasAsJava
+
 class LambdaTest extends AnyFlatSpec with Matchers with TestContainerUtils {
   override def config: Config = ConfigFactory.load()
 
   override def afterContainersStart(containers: containerDef.Container): Unit = super.afterContainersStart(containers)
 
+  val mockContext: Context = new Context {
+    override def getAwsRequestId: String = "testRequestId"
+    override def getLogGroupName: String = "testLogGroupName"
+    override def getLogStreamName: String = "testLogStreamName"
+    override def getFunctionName: String = "testFunctionName"
+    override def getFunctionVersion: String = "testFunctionVersion"
+    override def getInvokedFunctionArn: String = "testInvokedFunctionArn"
+    override def getIdentity: CognitoIdentity = ???
+    override def getClientContext: ClientContext = ???
+    override def getRemainingTimeInMillis: Int = ???
+    override def getMemoryLimitInMB: Int = ???
+    override def getLogger: LambdaLogger = ???
+  }
+
   "The Lambda class" should "return an APIGateWayResponseEvent with the correct number of references" in withContainers { case container: DynaliteContainer =>
     val client = createDynamoDbClient(container)
     val input = Lambda.Input(numberOfReferences = 3)
@@ -40,4 +57,32 @@ class LambdaTest extends AnyFlatSpec with Matchers with TestContainerUtils {
     actual.getStatusCode shouldBe expected.getStatusCode
     actual.getBody should include(expected.getBody)
   }
+
+  "The Lambda class" should "return an APIGateWayResponseEvent with body containing exception message if numberofrefs isn't an integer" in {
+    val lambda = new Lambda()
+    val queryParam = config.getString("dynamodb.queryParam")
+    val queryParams = Map(queryParam -> "abc").asJava
+    val event = new APIGatewayProxyRequestEvent()
+    event.setQueryStringParameters(queryParams)
+
+    val result = lambda.handleRequest(event, mockContext)
+    val expected: APIGatewayProxyResponseEvent = new APIGatewayProxyResponseEvent()
+      .withStatusCode(500)
+      .withBody(s"""For input string: "abc"""")
+    result shouldBe expected
+  }
+
+  "The Lambda class" should "return an APIGateWayResponseEvent with body containing exception message if numberofrefs exceeds the limit" in {
+    val lambda = new Lambda()
+    val queryParam = config.getString("dynamodb.queryParam")
+    val queryParams = Map(queryParam -> "5001").asJava
+    val event = new APIGatewayProxyRequestEvent()
+    event.setQueryStringParameters(queryParams)
+
+    val result = lambda.handleRequest(event, mockContext)
+    val expected: APIGatewayProxyResponseEvent = new APIGatewayProxyResponseEvent()
+      .withStatusCode(500)
+      .withBody(s"""$queryParam is greater than 5000""")
+    result shouldBe expected
+  }
 }

terraform/root_main.tf

@@ -43,6 +43,7 @@ module "reference_generator_lambda" {
     REFERENCE_KEY_VALUE = "fileCounter"
     REFERENCE_COUNTER   = "referenceCounter"
     QUERY_PARAM         = "numberofrefs"
+    REFERENCE_LIMIT     = 5000
   }
   runtime         = "java11"
   timeout_seconds = 60