ISSUE-133 Adding UID STORE command

Author: quantranhong1999

README.adoc

@@ -43,6 +43,7 @@ For now only the following commands are supported:
  - UID COPY
  - UID EXPUNGE
  - UID MOVE
+ - UID STORE
 
 == Supported checks
 

src/it/scala-2.12/com/linagora/gatling/imap/scenario/it/ImapAuthenticationScenarioIT.scala

@@ -2,7 +2,7 @@ package com.linagora.gatling.imap.scenario.it
 
 import com.linagora.gatling.imap.Fixture.bart
 import com.linagora.gatling.imap.PreDef.imap
-import com.linagora.gatling.imap.scenario.{ImapAuthenticationScenario, ImapCapabilityScenario, ImapCheckScenario, ImapCloseScenario, ImapCopyMessageScenario, ImapCreateFolderScenario, ImapDeleteFolderScenario, ImapEnableScenario, ImapExamineFolderScenario, ImapExpungeScenario, ImapGetAclScenario, ImapGetQuotaRootScenario, ImapGetQuotaScenario, ImapIdleScenario, ImapLogoutScenario, ImapLsubScenario, ImapMoveMessageScenario, ImapMyRightsScenario, ImapNamespaceScenario, ImapNoopScenario, ImapRenameFolderScenario, ImapSearchScenario, ImapSimpleScenario, ImapStatusScenario, ImapSubscribeScenario, ImapUIDFetchScenario, ImapUidCopyMessageScenario, ImapUidExpungeMessageScenario, ImapUidMoveMessageScenario, ImapUnselectScenario, ImapUnsubscribeScenario}
+import com.linagora.gatling.imap.scenario.{ImapAuthenticationScenario, ImapCapabilityScenario, ImapCheckScenario, ImapCloseScenario, ImapCopyMessageScenario, ImapCreateFolderScenario, ImapDeleteFolderScenario, ImapEnableScenario, ImapExamineFolderScenario, ImapExpungeScenario, ImapGetAclScenario, ImapGetQuotaRootScenario, ImapGetQuotaScenario, ImapIdleScenario, ImapLogoutScenario, ImapLsubScenario, ImapMoveMessageScenario, ImapMyRightsScenario, ImapNamespaceScenario, ImapNoopScenario, ImapRenameFolderScenario, ImapSearchScenario, ImapSimpleScenario, ImapStatusScenario, ImapSubscribeScenario, ImapUIDFetchScenario, ImapUidCopyMessageScenario, ImapUidExpungeMessageScenario, ImapUidMoveMessageScenario, ImapUidStoreScenario, ImapUnselectScenario, ImapUnsubscribeScenario}
 import com.linagora.gatling.imap.{CyrusServer, Fixture, JamesServer, RunningServer}
 import io.gatling.core.feeder.FeederBuilder
 import io.gatling.core.funspec.GatlingFunSpec
@@ -43,6 +43,9 @@ class ImapUIDFetchScenarioIT extends BaseIt(CyrusServer.start()) {
   scenario(ImapUIDFetchScenario(_))
 }
 
+class ImapUidStoreScenarioIT extends BaseIt(CyrusServer.start()) {
+  scenario(ImapUidStoreScenario(_))
+}
 
 class ImapAuthenticationScenarioJamesIT extends BaseIt(JamesServer.start()) {
   scenario(ImapAuthenticationScenario(_))

src/main/scala-2.12/com/linagora/gatling/imap/action/ImapActionBuilder.scala

@@ -118,6 +118,9 @@ class ImapActionBuilder(requestName: String) {
   def store(sequence: Expression[MessageRanges], flags: Expression[StoreFlags]): ImapStoreActionBuilder =
     ImapStoreActionBuilder(requestName, sequence, flags, Seq.empty)
 
+  def uidStore(sequence: Expression[MessageRanges], flags: Expression[StoreFlags]): ImapUidStoreActionBuilder =
+    ImapUidStoreActionBuilder(requestName, sequence, flags, Seq.empty)
+
   def status(mailbox: Expression[String], items: Expression[StatusItems]): ImapStatusActionBuilder =
     ImapStatusActionBuilder(requestName, mailbox, items, Seq.empty)
 
@@ -462,6 +465,15 @@ case class ImapStoreActionBuilder(requestName: String, sequence: Expression[Mess
   override val actionName: String = "store-action"
 }
 
+case class ImapUidStoreActionBuilder(requestName: String, sequence: Expression[MessageRanges], flags: Expression[StoreFlags], private val checks: Seq[ImapCheck]) extends ImapCommandActionBuilder {
+  def check(checks: ImapCheck*): ImapUidStoreActionBuilder = copy(checks = this.checks ++ checks)
+
+  override def props(ctx: ImapActionContext): Props =
+    UidStoreAction.props(ctx, requestName, checks, sequence, flags)
+
+  override val actionName: String = "uid-store-action"
+}
+
 case class ImapExpungeActionBuilder(requestName: String, private val checks: Seq[ImapCheck]) extends ImapCommandActionBuilder {
   def check(checks: ImapCheck*): ImapExpungeActionBuilder = copy(checks = this.checks ++ checks)
 

src/main/scala-2.12/com/linagora/gatling/imap/action/UidStoreAction.scala

@@ -0,0 +1,30 @@
+package com.linagora.gatling.imap.action
+
+import akka.actor.Props
+import com.linagora.gatling.imap.check.ImapCheck
+import com.linagora.gatling.imap.protocol.command.{MessageRanges, StoreFlags}
+import com.linagora.gatling.imap.protocol.{Command, UserId}
+import io.gatling.commons.validation.Validation
+import io.gatling.core.action.ValidatedActionActor
+import io.gatling.core.session._
+
+import scala.collection.immutable.Seq
+
+object UidStoreAction {
+  def props(imapContext: ImapActionContext, requestname: String, checks: Seq[ImapCheck], sequence: Expression[MessageRanges], flags: Expression[StoreFlags]) =
+    Props(new UidStoreAction(imapContext, requestname, checks, sequence, flags))
+}
+
+class UidStoreAction(val imapContext: ImapActionContext, val requestName: String, override val checks: Seq[ImapCheck], sequence: Expression[MessageRanges], flags: Expression[StoreFlags]) extends ValidatedActionActor with ImapActionActor {
+
+  override protected def executeOrFail(session: Session): Validation[_] = {
+    for {
+      sequence <- sequence(session)
+      flags <- flags(session)
+    } yield {
+      val id: Long = session.userId
+      val handler = handleResponse(session, imapContext.clock.nowMillis)
+      sessions.tell(Command.UidStore(UserId(id), sequence, flags), handler)
+    }
+  }
+}

src/main/scala-2.12/com/linagora/gatling/imap/protocol/ImapSessions.scala

@@ -192,6 +192,9 @@ private class ImapSession(client: => ImapAsyncClient, protocol: ImapProtocol) ex
     case cmd@Command.Store(_, _, _) =>
       val handler = context.actorOf(StoreHandler.props(session), genName("store"))
       handler forward cmd
+    case cmd@Command.UidStore(_, _, _) =>
+      val handler = context.actorOf(UidStoreHandler.props(session), genName("uidStore"))
+      handler forward cmd
     case cmd@Command.Expunge(_) =>
       val handler = context.actorOf(ExpungeHandler.props(session), genName("expunge"))
       handler forward cmd

src/main/scala-2.12/com/linagora/gatling/imap/protocol/command/UidStoreHandler.scala

@@ -0,0 +1,50 @@
+package com.linagora.gatling.imap.protocol.command
+
+import java.util.function.Consumer
+
+import akka.actor.{ActorRef, Props}
+import com.linagora.gatling.imap.protocol.{Command, ImapResponses, Response}
+import com.yahoo.imapnio.async.client.ImapAsyncSession
+import com.yahoo.imapnio.async.request.UidStoreFlagsCommand
+import com.yahoo.imapnio.async.response.ImapAsyncResponse
+import io.gatling.core.akka.BaseActor
+
+import scala.collection.immutable.Seq
+
+object UidStoreHandler {
+  def props(session: ImapAsyncSession) = Props(new UidStoreHandler(session))
+}
+
+class UidStoreHandler(session: ImapAsyncSession) extends BaseActor {
+  override def receive: Receive = {
+    case Command.UidStore(userId, sequence, flags) =>
+      logger.trace(s"UidStoreHandler for user : ${userId.value}, on actor ${self.path} responding to ${sender.path}")
+      context.become(waitCallback(sender()))
+
+      val responseCallback: Consumer[ImapAsyncResponse] = responses => {
+        import collection.JavaConverters._
+
+        val responsesList = ImapResponses(responses.getResponseLines.asScala.to[Seq])
+        logger.trace(s"On response for $userId :\n ${responsesList.mkString("\n")}")
+        self !  Response.UidStoreResult(responsesList)
+      }
+
+      val errorCallback: Consumer[Exception] = e => {
+        logger.trace(s"${getClass.getSimpleName} command failed", e)
+        logger.error(s"${getClass.getSimpleName} command failed")
+        sender ! e
+        context.stop(self)
+      }
+
+      val future = session.execute(new UidStoreFlagsCommand(sequence.asImap, flags.asImap, flags.action, true))
+      future.setDoneCallback(responseCallback)
+      future.setExceptionCallback(errorCallback)
+  }
+
+  def waitCallback(sender: ActorRef): Receive = {
+    case msg@Response.UidStoreResult(_) =>
+      logger.trace(s"UidStoreHandler respond to ${sender.path} with $msg")
+      sender ! msg
+      context.stop(self)
+  }
+}

src/main/scala-2.12/com/linagora/gatling/imap/protocol/package.scala

@@ -118,6 +118,8 @@ package object protocol {
     case class Status(userId: UserId, mailbox: String, items: StatusItems) extends Command
 
     case class Store(userId: UserId, sequence: MessageRanges, flags: StoreFlags) extends Command
+
+    case class UidStore(userId: UserId, sequence: MessageRanges, flags: StoreFlags) extends Command
   }
 
   sealed trait Response {
@@ -199,6 +201,8 @@ package object protocol {
 
     case class Stored(responses: ImapResponses) extends Response
 
+    case class UidStoreResult(responses: ImapResponses) extends Response
+
     case class Appended(responses: ImapResponses) extends Response
 
     case class Disconnected(message: String)

src/main/scala-2.12/com/linagora/gatling/imap/scenario/ImapUidStoreScenario.scala

@@ -0,0 +1,48 @@
+package com.linagora.gatling.imap.scenario
+
+import java.util.Calendar
+
+import com.linagora.gatling.imap.PreDef._
+import com.linagora.gatling.imap.protocol.command.FetchAttributes.AttributeList
+import com.linagora.gatling.imap.protocol.command.MessageRange._
+import com.linagora.gatling.imap.protocol.command.{MessageRanges, Silent, StoreFlags}
+import io.gatling.core.Predef._
+import io.gatling.core.feeder.FeederBuilder
+import io.gatling.core.scenario.Simulation
+import io.gatling.core.structure.ScenarioBuilder
+
+import scala.collection.immutable.Seq
+import scala.concurrent.duration._
+
+object ImapUidStoreScenario extends Simulation {
+  private val receiveEmail = exec(imap("append").append("INBOX", Some(Seq("\\Flagged")), Option.empty[Calendar],
+    """From: [email protected]
+      |To: [email protected]
+      |Subject: test subject
+      |
+      |Test content
+      |abcdefghijklmnopqrstuvwxyz
+      |0123456789""".stripMargin).check(ok))
+
+  def apply(feeder: FeederBuilder): ScenarioBuilder = scenario("Imap")
+    .feed(feeder)
+    .pause(1 second)
+    .exec(imap("Connect").connect()).exitHereIfFailed
+    .exec(imap("login").login("${username}", "${password}").check(ok))
+    .during(1 minute) {
+        exec(imap("list").list("", "*").check(ok, hasFolder("INBOX")))
+        .pause(200 milli)
+        .exec(imap("select").select("INBOX").check(ok))
+        .pause(200 milli)
+        .exec(receiveEmail)
+        .pause(200 milli)
+        .exec(imap("fetch").fetch(MessageRanges(Last()), AttributeList("BODY[HEADER]", "UID", "BODY[TEXT]")).check(ok))
+        .pause(200 milli)
+        .exec(imap("uidStore").uidStore(MessageRanges(Last()), StoreFlags.add(Silent.Disable(), "\\Deleted")).check(ok))
+        .pause(200 milli)
+        .exec(imap("expunge").expunge().check(ok))
+        .pause(200 milli)
+        .exec(imap("select check is empty").select("INBOX").check(hasNoExists))
+    }
+
+}