diff --git a/CHANGES.txt b/CHANGES.txt
index 001418f..437774d 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -122,6 +122,8 @@
  
 BUG FIXES
 
+PIG-5245: TestGrunt.testStopOnFailure is flaky (rohini)
+
 PIG-5244: Several unit tests are failing in Tez mode after merging spark branch (nkollar via szita)
 
 PIG-5207: BugFix e2e tests fail on spark (szita)
diff --git a/test/org/apache/pig/test/MiniCluster.java b/test/org/apache/pig/test/MiniCluster.java
index 17bb866..e0dd428 100644
--- a/test/org/apache/pig/test/MiniCluster.java
+++ b/test/org/apache/pig/test/MiniCluster.java
@@ -25,6 +25,7 @@
 import org.apache.hadoop.fs.FileSystem;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.hdfs.MiniDFSCluster;
+import org.apache.hadoop.mapreduce.MRJobConfig;
 import org.apache.hadoop.mapreduce.filecache.DistributedCache;
 import org.apache.hadoop.mapreduce.v2.MiniMRYarnCluster;
 import org.apache.pig.ExecType;
@@ -79,6 +80,7 @@
             Configuration config = new Configuration();
             config.set("yarn.scheduler.capacity.root.queues", "default");
             config.set("yarn.scheduler.capacity.root.default.capacity", "100");
+            config.set("yarn.scheduler.capacity.maximum-am-resource-percent", "0.1");
             m_dfs = new MiniDFSCluster(config, dataNodes, true, null);
             m_fileSys = m_dfs.getFileSystem();
             m_dfs_conf = m_dfs.getConfiguration(0);
@@ -97,8 +99,12 @@
             m_conf.set(FileSystem.FS_DEFAULT_NAME_KEY, m_dfs_conf.get(FileSystem.FS_DEFAULT_NAME_KEY));
             m_conf.unset(MRConfiguration.JOB_CACHE_FILES);
 
-            m_conf.setInt(MRConfiguration.IO_SORT_MB, 200);
-            m_conf.set(MRConfiguration.CHILD_JAVA_OPTS, "-Xmx512m");
+            m_conf.setInt(MRConfiguration.IO_SORT_MB, 50);
+            m_conf.set(MRConfiguration.CHILD_JAVA_OPTS, "-Xmx384m");
+            m_conf.setInt(MRJobConfig.MAP_MEMORY_MB, 512);
+            m_conf.setInt(MRJobConfig.REDUCE_MEMORY_MB, 512);
+            m_conf.set(MRJobConfig.MR_AM_COMMAND_OPTS, "-Xmx384m");
+            m_conf.setInt(MRJobConfig.MR_AM_VMEM_MB, 512);
 
             m_conf.setInt(MRConfiguration.SUMIT_REPLICATION, 2);
             m_conf.setInt(MRConfiguration.MAP_MAX_ATTEMPTS, 2);
diff --git a/test/org/apache/pig/test/TestGrunt.java b/test/org/apache/pig/test/TestGrunt.java
index bf4b5e2..ed84dba 100644
--- a/test/org/apache/pig/test/TestGrunt.java
+++ b/test/org/apache/pig/test/TestGrunt.java
@@ -981,6 +981,7 @@
 
     @Test
     public void testStopOnFailure() throws Throwable {
+        Assume.assumeFalse("Skip this test for Spark", Util.isSparkExecType(cluster.getExecType()));
         PigServer server = new PigServer(cluster.getExecType(), cluster.getProperties());
         PigContext context = server.getPigContext();
         context.getProperties().setProperty("stop.on.failure", ""+true);
diff --git a/test/org/apache/pig/test/TezMiniCluster.java b/test/org/apache/pig/test/TezMiniCluster.java
index 0d8871f..2e76d76 100644
--- a/test/org/apache/pig/test/TezMiniCluster.java
+++ b/test/org/apache/pig/test/TezMiniCluster.java
@@ -59,8 +59,6 @@
         try {
             // Write tez-site.xml
             Configuration tez_conf = new Configuration(false);
-            // TODO PIG-3659 - Remove this once memory management is fixed
-            tez_conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_IO_SORT_MB, "20");
             tez_conf.set(TezRuntimeConfiguration.TEZ_RUNTIME_OPTIMIZE_LOCAL_FETCH, "false");
             tez_conf.set("tez.lib.uris", "hdfs:///tez,hdfs:///tez/lib");
             // Set to a lower value so that tests don't get stuck for long because of 1 AM running at a time
diff --git a/test/org/apache/pig/test/YarnMiniCluster.java b/test/org/apache/pig/test/YarnMiniCluster.java
index 6645557..9dac04b 100644
--- a/test/org/apache/pig/test/YarnMiniCluster.java
+++ b/test/org/apache/pig/test/YarnMiniCluster.java
@@ -75,20 +75,33 @@
             }
             hdfs_site.writeXml(new FileOutputStream(HDFS_CONF_FILE));
 
+            m_dfs_conf.set("yarn.scheduler.capacity.root.queues", "default");
+            m_dfs_conf.set("yarn.scheduler.capacity.root.default.capacity", "100");
+            m_dfs_conf.set("yarn.scheduler.capacity.maximum-am-resource-percent", "0.1");
             // Build mini YARN cluster
             m_mr = new MiniMRYarnCluster("PigMiniCluster", 2);
             m_mr.init(m_dfs_conf);
             m_mr.start();
             m_mr_conf = m_mr.getConfig();
+
             File libDir = new File(System.getProperty("ivy.lib.dir", "build/ivy/lib/Pig"));
+            File sparkLibDir = new File(System.getProperty("ivy.lib.dir.spark", "build/ivy/lib/Pig/spark"));
             File classesDir = new File(System.getProperty("build.classes", "build/classes"));
             File testClassesDir = new File(System.getProperty("test.build.classes", "test/build/classes"));
+
             String classpath = libDir.getAbsolutePath() + "/*"
-                + File.pathSeparator + classesDir.getAbsolutePath()
-                + File.pathSeparator + testClassesDir.getAbsolutePath();
+                    + (Util.isSparkExecType(getExecType()) ? File.pathSeparator + sparkLibDir.getAbsolutePath() + "/*" : "")
+                    + File.pathSeparator + classesDir.getAbsolutePath()
+                    + File.pathSeparator + testClassesDir.getAbsolutePath();
+
             m_mr_conf.set(YarnConfiguration.YARN_APPLICATION_CLASSPATH, classpath);
-            m_mr_conf.set(MRJobConfig.MAP_JAVA_OPTS, "-Xmx512m");
-            m_mr_conf.set(MRJobConfig.REDUCE_JAVA_OPTS, "-Xmx512m");
+            m_mr_conf.setInt(MRJobConfig.IO_SORT_MB, 50);
+            m_mr_conf.set(MRJobConfig.MAP_JAVA_OPTS, "-Xmx384m");
+            m_mr_conf.set(MRJobConfig.REDUCE_JAVA_OPTS, "-Xmx384m");
+            m_mr_conf.setInt(MRJobConfig.MAP_MEMORY_MB, 512);
+            m_mr_conf.setInt(MRJobConfig.REDUCE_MEMORY_MB, 512);
+            m_mr_conf.set(MRJobConfig.MR_AM_COMMAND_OPTS, "-Xmx384m");
+            m_mr_conf.setInt(MRJobConfig.MR_AM_VMEM_MB, 512);
 
             Configuration mapred_site = new Configuration(false);
             Configuration yarn_site = new Configuration(false);
@@ -137,4 +150,4 @@
             m_mr = null;
         }
     }
-}
\ No newline at end of file
+}