From 946435f25d2c378cb1fdc802329c6640bfc4373e Mon Sep 17 00:00:00 2001
From: Alan Dipert <alan@dipert.org>
Date: Thu, 5 Mar 2020 18:11:42 +0000
Subject: [PATCH 1/7] Add class to ShinyMockSession and fix tests

---
 R/mock-session.R  | 1 -
 R/modules.R       | 4 +++-
 inst/_pkgdown.yml | 2 +-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/R/mock-session.R b/R/mock-session.R
index a88871cf30..a90b060f2e 100644
--- a/R/mock-session.R
+++ b/R/mock-session.R
@@ -80,7 +80,6 @@ extract <- function(promise) {
 MockShinySession <- R6Class(
   'MockShinySession',
   portable = FALSE,
-  class = FALSE,
   public = list(
     #' @field env The environment associated with the session.
     env = NULL,
diff --git a/R/modules.R b/R/modules.R
index 24bea6a346..a103ef97f8 100644
--- a/R/modules.R
+++ b/R/modules.R
@@ -128,7 +128,9 @@ moduleServer <- function(id, module, session = getDefaultReactiveDomain()) {
 #' @rdname moduleServer
 #' @export
 callModule <- function(module, id, ..., session = getDefaultReactiveDomain()) {
-  if (!inherits(session, "ShinySession") && !inherits(session, "session_proxy")) {
+  if (!inherits(session, "ShinySession")
+      && !inherits(session, "session_proxy")
+      && !inherits(session, "MockShinySession")) {
     stop("session must be a ShinySession or session_proxy object.")
   }
   childScope <- session$makeScope(id)
diff --git a/inst/_pkgdown.yml b/inst/_pkgdown.yml
index 49ecfe2b38..8112bf3ca3 100644
--- a/inst/_pkgdown.yml
+++ b/inst/_pkgdown.yml
@@ -207,7 +207,7 @@ reference:
     desc: Functions for modularizing Shiny apps
     contents:
       - NS
-      - callModule
+      - moduleServer
   - title: Embedding
     desc: Functions that are intended for third-party packages that embed Shiny applications.
     contents:

From 14b572e11530aa63e616000552c00c73c3213625 Mon Sep 17 00:00:00 2001
From: Alan Dipert <alan@dipert.org>
Date: Thu, 5 Mar 2020 19:25:33 +0000
Subject: [PATCH 2/7] Passing existing tests

---
 R/modules.R     | 37 +++++++++++++++++++++++++++++++++----
 R/test-module.R | 40 +++++++++++++++++++++++++++-------------
 2 files changed, 60 insertions(+), 17 deletions(-)

diff --git a/R/modules.R b/R/modules.R
index a103ef97f8..9b9db38638 100644
--- a/R/modules.R
+++ b/R/modules.R
@@ -125,21 +125,50 @@ moduleServer <- function(id, module, session = getDefaultReactiveDomain()) {
 }
 
 
+parentSession <- function(session) {
+  if (inherits(session, c("ShinySession", "MockShinySession")))
+    return(session)
+
+  if (!inherits(session, "session_proxy"))
+    stop("session must be a ShinySession, MockShinySession, or session_proxy object.")
+
+  while (inherits(session, "session_proxy"))
+    session <- session$parent
+
+  session
+}
+
 #' @rdname moduleServer
 #' @export
 callModule <- function(module, id, ..., session = getDefaultReactiveDomain()) {
-  if (!inherits(session, "ShinySession")
-      && !inherits(session, "session_proxy")
-      && !inherits(session, "MockShinySession")) {
-    stop("session must be a ShinySession or session_proxy object.")
+  if (!inherits(session, c("ShinySession", "MockShinySession", "session_proxy"))) {
+    stop("session must be a ShinySession, MockShinySession, or session_proxy object.")
   }
+
   childScope <- session$makeScope(id)
+  parent <- parentSession(session)
+  trace <- rlang::trace_back()
+
+  if (inherits(parent, "MockShinySession")
+      && (sys.nframe() >= 2)
+      && (as.character(sys.call(1)[[1]]) == "moduleServer")) {
+    # If the module is under test *and* was called by moduleServer(), modify the
+    # module function locally by inserting the equivalent of `session$env <-
+    # environment()` at the beginning of its body. A similar operation is
+    # performed by .testModule() if the module is *not* called through
+    # moduleServer() but is under test. See .testModule() for details.
+    body(module) <- rlang::expr({
+      base::assign("env", base::environment(), envir = !!parent)
+      !!!body(module)
+    })
+  }
 
   withReactiveDomain(childScope, {
     if (!is.function(module)) {
       stop("module argument must be a function")
     }
 
+    # TODO use rlang::execute to support dynamic dots?
     module(childScope$input, childScope$output, childScope, ...)
   })
 }
diff --git a/R/test-module.R b/R/test-module.R
index 0e17f05a35..76e4eb7f74 100644
--- a/R/test-module.R
+++ b/R/test-module.R
@@ -64,23 +64,37 @@ testModule <- function(module, expr, ...) {
   )
 }
 
+isOldModule <- function(func) {
+  stopifnot(is.function(func))
+  required <- c("input", "output", "session")
+  declared <- names(formals(func))
+  setequal(required, intersect(required, declared))
+}
+
 #' @noRd
-#' @importFrom withr with_options
 .testModule <- function(module, quosure, dots, env) {
-  # Modify the module function locally by inserting `session$env <-
-  # environment()` at the beginning of its body. The dynamic environment of the
-  # module function is saved so that it may be referenced after the module
-  # function has returned. The saved dynamic environment is the basis for the
-  # `data` argument of tidy_eval() when used below to evaluate `quosure`, the
-  # test code expression.
-  body(module) <- rlang::expr({
-    session$env <- base::environment()
-    !!!body(module)
-  })
-
   session <- MockShinySession$new()
   on.exit(if (!session$isClosed()) session$close())
-  args <- append(dots, list(input = session$input, output = session$output, session = session))
+
+  if (isOldModule(module)) {
+    # If the module is an "old-style" module that accepts input, output, and
+    # session parameters, modify the function locally by inserting `session$env
+    # <- environment()` at the beginning of its body. The dynamic environment of
+    # the module function is saved so that it may be referenced after the module
+    # function has returned. The saved dynamic environment is the basis for the
+    # `data` argument of tidy_eval() when used below to evaluate `quosure`.
+    body(module) <- rlang::expr({
+      session$env <- base::environment()
+      !!!body(module)
+    })
+    args <- append(dots, list(input = session$input, output = session$output, session = session))
+  } else {
+    # If the module is a "new-style" module, we rely on logic in callModule()
+    # that instruments the function if the session is a MockShinySession.
+    # Appending additional arguments is not necessary, as input/output/session
+    # will be provided in moduleServer().
+    args <- dots
+  }
 
   isolate(
     withReactiveDomain(

From 6fa332aa77fda444c7ecc881afc6b2493d40c338 Mon Sep 17 00:00:00 2001
From: Alan Dipert <alan@dipert.org>
Date: Thu, 5 Mar 2020 19:46:35 +0000
Subject: [PATCH 3/7] Add changes and a failing test

---
 R/modules.R                             |  3 +--
 R/test-module.R                         |  4 +++-
 tests/testthat/test-test-moduleServer.R | 23 +++++++++++++++++++++++
 3 files changed, 27 insertions(+), 3 deletions(-)
 create mode 100644 tests/testthat/test-test-moduleServer.R

diff --git a/R/modules.R b/R/modules.R
index 9b9db38638..bd61aea045 100644
--- a/R/modules.R
+++ b/R/modules.R
@@ -147,11 +147,10 @@ callModule <- function(module, id, ..., session = getDefaultReactiveDomain()) {
 
   childScope <- session$makeScope(id)
   parent <- parentSession(session)
-  trace <- rlang::trace_back()
 
   if (inherits(parent, "MockShinySession")
       && (sys.nframe() >= 2)
-      && (as.character(sys.call(1)[[1]]) == "moduleServer")) {
+      && (as.character(sys.call(sys.nframe() - 1)[[1]]) == "moduleServer")) {
     # If the module is under test *and* was called by moduleServer(), modify the
     # module function locally by inserting the equivalent of `session$env <-
     # environment()` at the beginning of its body. A similar operation is
diff --git a/R/test-module.R b/R/test-module.R
index 76e4eb7f74..c8bc4260e1 100644
--- a/R/test-module.R
+++ b/R/test-module.R
@@ -92,7 +92,8 @@ isOldModule <- function(func) {
     # If the module is a "new-style" module, we rely on logic in callModule()
     # that instruments the function if the session is a MockShinySession.
     # Appending additional arguments is not necessary, as input/output/session
-    # will be provided in moduleServer().
+    # will be provided in moduleServer(). `id` is also provided via
+    # moduleServer().
     args <- dots
   }
 
@@ -101,6 +102,7 @@ isOldModule <- function(func) {
       session,
       withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
         # Assigning to `$returned` causes a flush to happen automatically.
+        # TODO Wrong for new-style modules; fix
         session$returned <- do.call(module, args)
       })
     )
diff --git a/tests/testthat/test-test-moduleServer.R b/tests/testthat/test-test-moduleServer.R
new file mode 100644
index 0000000000..2b179e983b
--- /dev/null
+++ b/tests/testthat/test-test-moduleServer.R
@@ -0,0 +1,23 @@
+context("testModule-moduleServer")
+
+test_that("New-style modules work", {
+  counterServer <- local({
+    function(id) {
+      moduleServer(id, function(input, output, session) {
+        count <- reactiveVal(0)
+        observeEvent(input$button, {
+          count(count() + 1)
+        })
+        output$out <- renderText({
+          count()
+        })
+        count
+      })
+    }
+  })
+  testModule(counterServer, {
+    input$setInputs(button = 0)
+    input$setInputs(button = 1)
+    expect_equal(count(), 1)
+  }, id = "foob")
+})

From ed93d42a6eac556901020a1e223dadbe0adafab5 Mon Sep 17 00:00:00 2001
From: Alan Dipert <alan@dipert.org>
Date: Thu, 5 Mar 2020 21:24:38 +0000
Subject: [PATCH 4/7] Simplify differentiation strategy

---
 R/mock-session.R | 3 ++-
 R/modules.R      | 6 +++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/R/mock-session.R b/R/mock-session.R
index a90b060f2e..eb5234d194 100644
--- a/R/mock-session.R
+++ b/R/mock-session.R
@@ -100,7 +100,8 @@ MockShinySession <- R6Class(
     userData = NULL,
     #' @field progressStack A stack of progress objects
     progressStack = 'Stack',
-
+    #' @field TRUE when a moduleServer()-based module is under test
+    isModuleServer = FALSE,
     #' @description Create a new MockShinySession
     initialize = function() {
       private$.input <- ReactiveValues$new(dedupe = FALSE, label = "input")
diff --git a/R/modules.R b/R/modules.R
index bd61aea045..832caca9c7 100644
--- a/R/modules.R
+++ b/R/modules.R
@@ -121,6 +121,8 @@ createSessionProxy <- function(parentSession, ...) {
 #'
 #' @export
 moduleServer <- function(id, module, session = getDefaultReactiveDomain()) {
+  parent <- parentSession(session)
+  if (inherits(parent, c("MockShinySession"))) session$isModuleServer <- TRUE
   callModule(module, id, session = session)
 }
 
@@ -148,9 +150,7 @@ callModule <- function(module, id, ..., session = getDefaultReactiveDomain()) {
   childScope <- session$makeScope(id)
   parent <- parentSession(session)
 
-  if (inherits(parent, "MockShinySession")
-      && (sys.nframe() >= 2)
-      && (as.character(sys.call(sys.nframe() - 1)[[1]]) == "moduleServer")) {
+  if (inherits(parent, "MockShinySession") && parent$isModuleServer) {
     # If the module is under test *and* was called by moduleServer(), modify the
     # module function locally by inserting the equivalent of `session$env <-
     # environment()` at the beginning of its body. A similar operation is

From c2c0a0d8360648c490bac2c03f71856ecc5a0986 Mon Sep 17 00:00:00 2001
From: Alan Dipert <alan@dipert.org>
Date: Fri, 6 Mar 2020 00:35:35 +0000
Subject: [PATCH 5/7] A little churn

---
 R/modules.R     | 47 +++++++++++++++++------------------------------
 R/test-module.R | 30 ++++++++++++++++++------------
 2 files changed, 35 insertions(+), 42 deletions(-)

diff --git a/R/modules.R b/R/modules.R
index 832caca9c7..3a333b7708 100644
--- a/R/modules.R
+++ b/R/modules.R
@@ -121,25 +121,10 @@ createSessionProxy <- function(parentSession, ...) {
 #'
 #' @export
 moduleServer <- function(id, module, session = getDefaultReactiveDomain()) {
-  parent <- parentSession(session)
-  if (inherits(parent, c("MockShinySession"))) session$isModuleServer <- TRUE
+  if (inherits(session, c("MockShinySession"))) session$isModuleServer <- TRUE
   callModule(module, id, session = session)
 }
 
-
-parentSession <- function(session) {
-  if (inherits(session, c("ShinySession", "MockShinySession")))
-    return(session)
-
-  if (!inherits(session, "session_proxy"))
-    stop("session must be a ShinySession, MockShinySession, or session_proxy object.")
-
-  while (inherits(session, "session_proxy"))
-    session <- session$parent
-
-  session
-}
-
 #' @rdname moduleServer
 #' @export
 callModule <- function(module, id, ..., session = getDefaultReactiveDomain()) {
@@ -147,27 +132,29 @@ callModule <- function(module, id, ..., session = getDefaultReactiveDomain()) {
     stop("session must be a ShinySession, MockShinySession, or session_proxy object.")
   }
 
-  childScope <- session$makeScope(id)
-  parent <- parentSession(session)
-
-  if (inherits(parent, "MockShinySession") && parent$isModuleServer) {
+  if (inherits(session, "MockShinySession") && session$isModuleServer) {
     # If the module is under test *and* was called by moduleServer(), modify the
     # module function locally by inserting the equivalent of `session$env <-
     # environment()` at the beginning of its body. A similar operation is
     # performed by .testModule() if the module is *not* called through
     # moduleServer() but is under test. See .testModule() for details.
     body(module) <- rlang::expr({
-      base::assign("env", base::environment(), envir = !!parent)
+      base::assign("env", base::environment(), envir = !!session)
       !!!body(module)
     })
+    isolate(
+      withReactiveDomain(
+        session,
+        withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
+          session$returned <- module(session$input, session$output, session, ...)
+        })
+      )
+    )
+  } else {
+    childScope <- session$makeScope(id)
+    withReactiveDomain(childScope, {
+      if (!is.function(module)) stop("module argument must be a function")
+      module(childScope$input, childScope$output, childScope, ...)
+    })
   }
-
-  withReactiveDomain(childScope, {
-    if (!is.function(module)) {
-      stop("module argument must be a function")
-    }
-
-    # TODO use rlang::execute to support dynamic dots?
-    module(childScope$input, childScope$output, childScope, ...)
-  })
 }
diff --git a/R/test-module.R b/R/test-module.R
index c8bc4260e1..75924818e6 100644
--- a/R/test-module.R
+++ b/R/test-module.R
@@ -88,25 +88,31 @@ isOldModule <- function(func) {
       !!!body(module)
     })
     args <- append(dots, list(input = session$input, output = session$output, session = session))
+    isolate(
+      withReactiveDomain(
+        session,
+        withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
+          # Assigning to `$returned` causes a flush to happen automatically.
+          session$returned <- do.call(module, args)
+        })
+      )
+    )
   } else {
     # If the module is a "new-style" module, we rely on logic in callModule()
     # that instruments the function if the session is a MockShinySession.
     # Appending additional arguments is not necessary, as input/output/session
     # will be provided in moduleServer(). `id` is also provided via
     # moduleServer().
-    args <- dots
-  }
-
-  isolate(
-    withReactiveDomain(
-      session,
-      withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
-        # Assigning to `$returned` causes a flush to happen automatically.
-        # TODO Wrong for new-style modules; fix
-        session$returned <- do.call(module, args)
-      })
+    isolate(
+      withReactiveDomain(
+        session,
+        withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
+          # TODO Implement session$returned for new style modules
+          do.call(module, dots)
+        })
+      )
     )
-  )
+  }
 
   # Evaluate `quosure` in a reactive context, and in the provided `env`, but
   # with `env` masked by a shallow view of `session$env`, the environment that

From 2c2ca4b58ea1ca9197edc3d144ecc1eb95762f17 Mon Sep 17 00:00:00 2001
From: Alan Dipert <alan@dipert.org>
Date: Fri, 6 Mar 2020 23:40:27 +0000
Subject: [PATCH 6/7] simplifications

---
 R/mock-session.R | 21 +++++++++++++++-
 R/modules.R      | 52 +++++++++++++++++++--------------------
 R/test-module.R  | 63 +++++++++++-------------------------------------
 3 files changed, 60 insertions(+), 76 deletions(-)

diff --git a/R/mock-session.R b/R/mock-session.R
index eb5234d194..6a79c5df7d 100644
--- a/R/mock-session.R
+++ b/R/mock-session.R
@@ -69,6 +69,20 @@ extract <- function(promise) {
   stop("Single-bracket indexing of mockclientdata is not allowed.")
 }
 
+#' @noRd
+patchModuleFunction <- function(module) {
+  body(module) <- rlang::expr({
+    withr::with_options(base::list(`shiny.allowoutputreads` = TRUE), {
+      session$env <- base::environment()
+      session$returned <- {
+        !!!body(module)
+      }
+      session$returned
+    })
+  })
+  module
+}
+
 #' Mock Shiny Session
 #'
 #' @description
@@ -389,7 +403,12 @@ MockShinySession <- R6Class(
         self,
         input = .createReactiveValues(private$.input, readonly = TRUE, ns = ns),
         output = structure(.createOutputWriter(self, ns = ns), class = "shinyoutput"),
-        makeScope = function(namespace) self$makeScope(ns(namespace))
+        makeScope = function(namespace) self$makeScope(ns(namespace)),
+        setInputs = function(...) {
+          args <- list(...)
+          names(args) <- ns(names(args))
+          do.call(self$setInputs, args)
+        }
       )
     }
   ),
diff --git a/R/modules.R b/R/modules.R
index 3a333b7708..ca7002aaeb 100644
--- a/R/modules.R
+++ b/R/modules.R
@@ -121,40 +121,40 @@ createSessionProxy <- function(parentSession, ...) {
 #'
 #' @export
 moduleServer <- function(id, module, session = getDefaultReactiveDomain()) {
-  if (inherits(session, c("MockShinySession"))) session$isModuleServer <- TRUE
+  if (inherits(sessionFor(session), "MockShinySession")) {
+    # session is either a MockShinySession or a proxy for one
+    module <- patchModuleFunction(module)
+  }
   callModule(module, id, session = session)
 }
 
+#' @noRd
+sessionFor <- function(session) {
+  if (inherits(session, c("MockShinySession", "ShinySession")))
+    return(session)
+
+  if (!inherits(session, "session_proxy"))
+    stop("session must be a ShinySession, MockShinySession, or session_proxy object.")
+
+  while (inherits(session, "session_proxy"))
+    session <- session$parent
+
+  session
+}
+
 #' @rdname moduleServer
 #' @export
 callModule <- function(module, id, ..., session = getDefaultReactiveDomain()) {
   if (!inherits(session, c("ShinySession", "MockShinySession", "session_proxy"))) {
     stop("session must be a ShinySession, MockShinySession, or session_proxy object.")
   }
+  childScope <- session$makeScope(id)
 
-  if (inherits(session, "MockShinySession") && session$isModuleServer) {
-    # If the module is under test *and* was called by moduleServer(), modify the
-    # module function locally by inserting the equivalent of `session$env <-
-    # environment()` at the beginning of its body. A similar operation is
-    # performed by .testModule() if the module is *not* called through
-    # moduleServer() but is under test. See .testModule() for details.
-    body(module) <- rlang::expr({
-      base::assign("env", base::environment(), envir = !!session)
-      !!!body(module)
-    })
-    isolate(
-      withReactiveDomain(
-        session,
-        withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
-          session$returned <- module(session$input, session$output, session, ...)
-        })
-      )
-    )
-  } else {
-    childScope <- session$makeScope(id)
-    withReactiveDomain(childScope, {
-      if (!is.function(module)) stop("module argument must be a function")
-      module(childScope$input, childScope$output, childScope, ...)
-    })
-  }
+  withReactiveDomain(childScope, {
+    if (!is.function(module)) {
+      stop("module argument must be a function")
+    }
+
+    module(childScope$input, childScope$output, childScope, ...)
+  })
 }
diff --git a/R/test-module.R b/R/test-module.R
index 75924818e6..969ced4032 100644
--- a/R/test-module.R
+++ b/R/test-module.R
@@ -77,59 +77,24 @@ isOldModule <- function(func) {
   on.exit(if (!session$isClosed()) session$close())
 
   if (isOldModule(module)) {
-    # If the module is an "old-style" module that accepts input, output, and
-    # session parameters, modify the function locally by inserting `session$env
-    # <- environment()` at the beginning of its body. The dynamic environment of
-    # the module function is saved so that it may be referenced after the module
-    # function has returned. The saved dynamic environment is the basis for the
-    # `data` argument of tidy_eval() when used below to evaluate `quosure`.
-    body(module) <- rlang::expr({
-      session$env <- base::environment()
-      !!!body(module)
-    })
+    module <- patchModuleFunction(module)
     args <- append(dots, list(input = session$input, output = session$output, session = session))
-    isolate(
-      withReactiveDomain(
-        session,
-        withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
-          # Assigning to `$returned` causes a flush to happen automatically.
-          session$returned <- do.call(module, args)
-        })
-      )
-    )
   } else {
-    # If the module is a "new-style" module, we rely on logic in callModule()
-    # that instruments the function if the session is a MockShinySession.
-    # Appending additional arguments is not necessary, as input/output/session
-    # will be provided in moduleServer(). `id` is also provided via
-    # moduleServer().
-    isolate(
-      withReactiveDomain(
-        session,
-        withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
-          # TODO Implement session$returned for new style modules
-          do.call(module, dots)
-        })
-      )
-    )
+    args <- dots
   }
 
-  # Evaluate `quosure` in a reactive context, and in the provided `env`, but
-  # with `env` masked by a shallow view of `session$env`, the environment that
-  # was saved when the module function was invoked. flush is not needed before
-  # entering the loop because the first expr executed is `{`.
-  isolate({
-    withReactiveDomain(
-      session,
-      withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
-        rlang::eval_tidy(
-          quosure,
-          data = rlang::as_data_mask(as.list(session$env)),
-          env = env
-        )
-      })
-    )
-  })
+  withReactiveDomain(session, do.call(module, args))
+
+  withReactiveDomain(
+    session,
+    withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
+      rlang::eval_tidy(
+        quosure,
+        data = rlang::as_data_mask(as.list(session$env)),
+        env = env
+      )
+    })
+  )
 }
 
 #' Test an app's server-side logic

From a036aa46075271bf7c5f4c4d18480c39b9c01415 Mon Sep 17 00:00:00 2001
From: Alan Dipert <alan@dipert.org>
Date: Sat, 7 Mar 2020 00:10:07 +0000
Subject: [PATCH 7/7] still broken, progress on new modules and proxied
 mocksession

---
 R/mock-session.R | 26 +++++++++++++++++++++-----
 R/modules.R      |  5 +++--
 R/test-module.R  | 26 ++++++++++++++------------
 3 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/R/mock-session.R b/R/mock-session.R
index 6a79c5df7d..273da3e6f1 100644
--- a/R/mock-session.R
+++ b/R/mock-session.R
@@ -73,11 +73,10 @@ extract <- function(promise) {
 patchModuleFunction <- function(module) {
   body(module) <- rlang::expr({
     withr::with_options(base::list(`shiny.allowoutputreads` = TRUE), {
-      session$env <- base::environment()
-      session$returned <- {
+      session$setEnv(base::environment())
+      session$setReturned({
         !!!body(module)
-      }
-      session$returned
+      })
     })
   })
   module
@@ -395,21 +394,38 @@ MockShinySession <- R6Class(
     flushReact = function(){
       private$flush()
     },
+    setEnv = function(env) {
+      self$env <- env
+    },
+    setReturned = function(value) {
+      private$returnedVal <- value
+      private$flush()
+      value
+    },
     #' @description Create and return a namespace-specific session proxy.
     #' @param namespace Character vector indicating a namespace.
     makeScope = function(namespace) {
       ns <- NS(namespace)
-      createSessionProxy(
+      proxy <- createSessionProxy(
         self,
         input = .createReactiveValues(private$.input, readonly = TRUE, ns = ns),
         output = structure(.createOutputWriter(self, ns = ns), class = "shinyoutput"),
         makeScope = function(namespace) self$makeScope(ns(namespace)),
+        env = NULL,
+        returned = NULL,
+        setEnv = function(env) assign("env", env, envir = proxy),
+        setReturned = function(value) {
+          assign("returned", value, envir = proxy)
+          private$flush()
+          value
+        },
         setInputs = function(...) {
           args <- list(...)
           names(args) <- ns(names(args))
           do.call(self$setInputs, args)
         }
       )
+      proxy
     }
   ),
   private = list(
diff --git a/R/modules.R b/R/modules.R
index ca7002aaeb..4df36342c5 100644
--- a/R/modules.R
+++ b/R/modules.R
@@ -122,10 +122,11 @@ createSessionProxy <- function(parentSession, ...) {
 #' @export
 moduleServer <- function(id, module, session = getDefaultReactiveDomain()) {
   if (inherits(sessionFor(session), "MockShinySession")) {
-    # session is either a MockShinySession or a proxy for one
     module <- patchModuleFunction(module)
+    isolate(callModule(module, id, session = session))
+  } else {
+    callModule(module, id, session = session)
   }
-  callModule(module, id, session = session)
 }
 
 #' @noRd
diff --git a/R/test-module.R b/R/test-module.R
index 969ced4032..d4fabffae5 100644
--- a/R/test-module.R
+++ b/R/test-module.R
@@ -83,18 +83,20 @@ isOldModule <- function(func) {
     args <- dots
   }
 
-  withReactiveDomain(session, do.call(module, args))
-
-  withReactiveDomain(
-    session,
-    withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
-      rlang::eval_tidy(
-        quosure,
-        data = rlang::as_data_mask(as.list(session$env)),
-        env = env
-      )
-    })
-  )
+  isolate(withReactiveDomain(session, do.call(module, args)))
+
+  isolate({
+    withReactiveDomain(
+      session,
+      withr::with_options(list(`shiny.allowoutputreads`=TRUE), {
+        rlang::eval_tidy(
+          quosure,
+          data = rlang::as_data_mask(as.list(session$env)),
+          env = env
+        )
+      })
+    )
+  })
 }
 
 #' Test an app's server-side logic