diff --git a/src/goto-programs/replace_java_nondet.cpp b/src/goto-programs/replace_java_nondet.cpp
index 431acc0efb2..19adc917fbe 100644
--- a/src/goto-programs/replace_java_nondet.cpp
+++ b/src/goto-programs/replace_java_nondet.cpp
@@ -143,7 +143,19 @@ static bool is_assignment_from(
 
 /// Given an iterator into a list of instructions, modify the list to replace
 /// 'nondet' library functions with CBMC-native nondet expressions, and return
-/// an iterator to the next instruction to check.
+/// an iterator to the next instruction to check. It's important to note that
+/// this method also deals with the fact that in the GOTO program it assigns
+/// function return values to temporary variables, performs validation and then
+/// assigns the value into the actual variable.
+///
+/// Example:
+///
+///   return_tmp0=org.cprover.CProver.nondetWithoutNull:()Ljava/lang/Object;();
+///   ... Various validations of type and value here.
+///   obj = (<type-of-obj>)return_tmp0;
+///
+/// We're going to replace all of these lines with
+///   return_tmp0 = NONDET(<type-of-obj>)
 /// \param goto_program: The goto program to modify.
 /// \param target: A single step of the goto program which may be erased and
 ///   replaced.
@@ -161,40 +173,68 @@ static goto_programt::targett check_and_replace_target(
     return next_instr;
   }
 
-  // Look at the next instruction, ensure that it is an assignment
-  assert(next_instr->is_assign());
-  // Get the name of the LHS of the assignment
-  const auto &next_instr_assign_lhs=to_code_assign(next_instr->code).lhs();
-  if(!(next_instr_assign_lhs.id()==ID_symbol &&
-       !next_instr_assign_lhs.has_operands()))
+  // If we haven't removed returns yet, our function call will have a variable
+  // on its left hand side.
+  const bool remove_returns_not_run =
+    to_code_function_call(target->code).lhs().is_not_nil();
+
+  irep_idt return_identifier;
+  if(remove_returns_not_run)
   {
-    return next_instr;
+    return_identifier =
+      to_symbol_expr(to_code_function_call(target->code).lhs())
+        .get_identifier();
   }
-  const auto return_identifier=
-    to_symbol_expr(next_instr_assign_lhs).get_identifier();
+  else
+  {
+    // If not, we need to look at the next line instead.
+    DATA_INVARIANT(
+      next_instr->is_assign(),
+      "the code_function_callt for a nondet-returning library function should "
+      "either have a variable for the return value in its lhs() or the next "
+      "instruction should be an assignment of the return value to a temporary "
+      "variable");
+    const exprt &return_value_assignment =
+      to_code_assign(next_instr->code).lhs();
+
+    // If the assignment is null, return.
+    if(
+      !(return_value_assignment.id() == ID_symbol &&
+        !return_value_assignment.has_operands()))
+    {
+      return next_instr;
+    }
 
-  auto &instructions=goto_program.instructions;
-  const auto end=instructions.end();
+    // Otherwise it's the temporary variable.
+    return_identifier =
+      to_symbol_expr(return_value_assignment).get_identifier();
+  }
 
-  // Look for an instruction where this name is on the RHS of an assignment
-  const auto matching_assignment=std::find_if(
+  // Look for the assignment of the temporary return variable into our target
+  // variable.
+  const auto end = goto_program.instructions.end();
+  auto assignment_instruction = std::find_if(
     next_instr,
     end,
-    [&return_identifier](const goto_programt::instructiont &instr)
-    {
+    [&return_identifier](const goto_programt::instructiont &instr) {
       return is_assignment_from(instr, return_identifier);
     });
 
-  assert(matching_assignment!=end);
+  INVARIANT(
+    assignment_instruction != end,
+    "failed to find assignment of the temporary return variable into our "
+    "target variable");
 
   // Assume that the LHS of *this* assignment is the actual nondet variable
-  const auto &code_assign=to_code_assign(matching_assignment->code);
+  const auto &code_assign = to_code_assign(assignment_instruction->code);
   const auto nondet_var=code_assign.lhs();
   const auto source_loc=target->source_location;
 
   // Erase from the nondet function call to the assignment
-  const auto after_matching_assignment=std::next(matching_assignment);
-  assert(after_matching_assignment!=end);
+  const auto after_matching_assignment = std::next(assignment_instruction);
+  INVARIANT(
+    after_matching_assignment != end,
+    "goto_program missing END_FUNCTION instruction");
 
   std::for_each(
     target,
diff --git a/unit/java_bytecode/java_replace_nondet/Main.class b/unit/java_bytecode/java_replace_nondet/Main.class
new file mode 100644
index 00000000000..068fdf03f68
Binary files /dev/null and b/unit/java_bytecode/java_replace_nondet/Main.class differ
diff --git a/unit/java_bytecode/java_replace_nondet/replace_nondet.cpp b/unit/java_bytecode/java_replace_nondet/replace_nondet.cpp
new file mode 100644
index 00000000000..d78dba970fa
--- /dev/null
+++ b/unit/java_bytecode/java_replace_nondet/replace_nondet.cpp
@@ -0,0 +1,122 @@
+/*******************************************************************\
+
+ Module: Tests for the replace nondet pass to make sure it works both
+         when remove_returns has been run before and after the call.
+
+ Author: Diffblue Ltd.
+
+\*******************************************************************/
+
+#include <testing-utils/catch.hpp>
+#include <testing-utils/load_java_class.h>
+
+#include <goto-programs/goto_convert_functions.h>
+#include <goto-programs/remove_virtual_functions.h>
+#include <util/config.h>
+#include <goto-instrument/cover.h>
+#include <util/options.h>
+#include <iostream>
+#include <goto-programs/remove_returns.h>
+#include <goto-programs/replace_java_nondet.h>
+#include <goto-programs/remove_instanceof.h>
+
+void validate_method_removal(
+  std::list<goto_programt::instructiont> instructions)
+{
+  bool method_removed = true, replacement_nondet_exists = false;
+
+  // Quick loop to check that our method call has been replaced.
+  for(const auto &inst : instructions)
+  {
+    if(inst.is_assign())
+    {
+      const code_assignt &assignment = to_code_assign(inst.code);
+      if(assignment.rhs().id() == ID_side_effect)
+      {
+        const side_effect_exprt &see = to_side_effect_expr(assignment.rhs());
+        if(see.get_statement() == ID_nondet)
+        {
+          replacement_nondet_exists = true;
+        }
+      }
+    }
+
+    if(inst.is_function_call())
+    {
+      const code_function_callt &function_call =
+        to_code_function_call(inst.code);
+
+      // Small check to make sure the instruction is a symbol.
+      if(function_call.function().id() != ID_symbol)
+        continue;
+
+      const irep_idt function_id =
+        to_symbol_expr(function_call.function()).get_identifier();
+      if(
+        function_id ==
+        "java::org.cprover.CProver.nondetWithoutNull:()"
+        "Ljava/lang/Object;")
+      {
+        method_removed = false;
+      }
+    }
+  }
+
+  REQUIRE(method_removed);
+  REQUIRE(replacement_nondet_exists);
+}
+
+TEST_CASE(
+  "Load class with a generated java nondet method call, run remove returns "
+  "both before and after the nondet statements have been removed, check "
+  "results are as expected.",
+  "[core][java_bytecode][replace_nondet]")
+{
+  GIVEN("A class with a call to CProver.nondetWithoutNull()")
+  {
+    symbol_tablet raw_symbol_table = load_java_class(
+      "Main", "./java_bytecode/java_replace_nondet", "Main.replaceNondet");
+
+    journalling_symbol_tablet symbol_table =
+      journalling_symbol_tablet::wrap(raw_symbol_table);
+
+    null_message_handlert null_output;
+    goto_functionst functions;
+    goto_convert(symbol_table, functions, null_output);
+
+    const std::string function_name = "java::Main.replaceNondet:()V";
+    goto_functionst::goto_functiont &goto_function =
+      functions.function_map.at(function_name);
+
+    goto_model_functiont model_function(
+      symbol_table, functions, function_name, goto_function);
+
+    remove_instanceof(goto_function, symbol_table);
+
+    remove_virtual_functions(model_function);
+
+    WHEN("Remove returns is called before java nondet.")
+    {
+      remove_returns(model_function, [](const irep_idt &) { return false; });
+
+      replace_java_nondet(functions);
+
+      THEN("The nondet method call should have been removed.")
+      {
+        validate_method_removal(goto_function.body.instructions);
+      }
+    }
+
+    WHEN("Remove returns is called after java nondet.")
+    {
+      replace_java_nondet(functions);
+
+      remove_returns(model_function, [](const irep_idt &) { return false; });
+
+      THEN("The nondet method call should have been removed.")
+      {
+        validate_method_removal(goto_function.body.instructions);
+      }
+    }
+  }
+}
diff --git a/unit/java_bytecode/java_replace_nondet/test.java b/unit/java_bytecode/java_replace_nondet/test.java
new file mode 100644
index 00000000000..8a2a1a6d14b
--- /dev/null
+++ b/unit/java_bytecode/java_replace_nondet/test.java
@@ -0,0 +1,5 @@
+public class Main {
+    public void replaceNondet() {
+        Object test = org.cprover.CProver.nondetWithoutNull();
+    }
+}