diff --git a/wrangler-core/src/test/java/io/cdap/MockRelationalTransformContext.java b/wrangler-core/src/test/java/io/cdap/MockRelationalTransformContext.java
index 295cdb43c..8fb066802 100644
--- a/wrangler-core/src/test/java/io/cdap/MockRelationalTransformContext.java
+++ b/wrangler-core/src/test/java/io/cdap/MockRelationalTransformContext.java
@@ -28,9 +28,6 @@
 import java.util.Map;
 import java.util.Set;
 
-/**
- * Basic implementation of the {@link RelationalTranformContext} with single output port
- */
 public class MockRelationalTransformContext implements RelationalTranformContext {
     private Engine engine;
     private Map<String, Relation> inputMap;
@@ -39,7 +36,6 @@ public class MockRelationalTransformContext implements RelationalTranformContext
     private  FeatureFlagsProvider featureFlagsProvider;
     private Relation outputRelation;
 
-
     public MockRelationalTransformContext(Engine engine,
                                           Map<String, Relation> inputMap,
                                           Map<String, Schema> inputSchemas,
@@ -92,9 +88,6 @@ public boolean isFeatureEnabled(String name) {
         return this.featureFlagsProvider.isFeatureEnabled(name);
     }
 
-    public Relation getOutputRelation() {
-        return outputRelation;
-    }
 
     public Collection<Capability> getDefaultLanguageCapabilityList() {
         return Collections.singleton(StandardSQLCapabilities.POSTGRES);
diff --git a/wrangler-core/src/test/java/io/cdap/directives/column/RenameTest.java b/wrangler-core/src/test/java/io/cdap/directives/column/RenameTest.java
index 02691d3ee..169506adc 100644
--- a/wrangler-core/src/test/java/io/cdap/directives/column/RenameTest.java
+++ b/wrangler-core/src/test/java/io/cdap/directives/column/RenameTest.java
@@ -16,15 +16,7 @@
 
 package io.cdap.directives.column;
 
-import io.cdap.MockEngine;
-import io.cdap.MockExpression;
-import io.cdap.MockRelation;
-import io.cdap.MockRelationalTransformContext;
-import io.cdap.cdap.etl.api.relational.Engine;
-import io.cdap.cdap.etl.api.relational.Relation;
-import io.cdap.cdap.etl.api.relational.RelationalTranformContext;
 import io.cdap.wrangler.TestingRig;
-import io.cdap.wrangler.api.DirectiveParseException;
 import io.cdap.wrangler.api.RecipeException;
 import io.cdap.wrangler.api.Row;
 import org.junit.Assert;
@@ -32,10 +24,6 @@
 
 import java.util.Arrays;
 import java.util.List;
-
-import static io.cdap.RelationalDirectiveTest.runTransform;
-import static org.mockito.ArgumentMatchers.any;
-
 /**
  * Tests {@link Rename}
  */
diff --git a/wrangler-core/src/test/java/io/cdap/directives/transformation/RightTrimTest.java b/wrangler-core/src/test/java/io/cdap/directives/transformation/RightTrimTest.java
index 43b06c121..993eb6a55 100644
--- a/wrangler-core/src/test/java/io/cdap/directives/transformation/RightTrimTest.java
+++ b/wrangler-core/src/test/java/io/cdap/directives/transformation/RightTrimTest.java
@@ -90,12 +90,8 @@ public void testSentenceRightTrim() throws Exception {
   public void testRelationColumn() throws DirectiveParseException, RecipeException {
     MockRelation relation = new MockRelation(null, null);
     Engine engine = new MockEngine();
-    RelationalTranformContext relationalTranformContext = new MockRelationalTransformContext(engine, null
-            , null
-            ,
-            null
-            , null
-    );
+    RelationalTranformContext relationalTranformContext = new MockRelationalTransformContext(engine,
+            null, null, null, null);
     String[] recipe = {"rtrim testColumn"};
     Relation relation1 = runTransform(recipe, relationalTranformContext, relation);
     Assert.assertEquals(((MockRelation) relation1).getColumn(), "testColumn");
@@ -105,12 +101,8 @@ public void testRelationColumn() throws DirectiveParseException, RecipeException
   public void testRelationExpression() throws DirectiveParseException, RecipeException {
     MockRelation relation = new MockRelation(null, null);
     Engine engine = new MockEngine();
-    RelationalTranformContext relationalTranformContext = new MockRelationalTransformContext(engine, null
-            , null
-            ,
-            null
-            , null
-    );
+    RelationalTranformContext relationalTranformContext = new MockRelationalTransformContext(engine,
+            null, null, null, null);
     String[] recipe = {"rtrim testColumn"};
     Relation relation1 = runTransform(recipe, relationalTranformContext, relation);
     Assert.assertEquals(((MockExpression) ((MockRelation) relation1).getExpression()).getExpression(),
diff --git a/wrangler-transform/src/test/java/io/cdap/wrangler/WranglerTransformTest.java b/wrangler-transform/src/test/java/io/cdap/wrangler/WranglerTransformTest.java
index a40bd3cc7..f843931c4 100644
--- a/wrangler-transform/src/test/java/io/cdap/wrangler/WranglerTransformTest.java
+++ b/wrangler-transform/src/test/java/io/cdap/wrangler/WranglerTransformTest.java
@@ -16,43 +16,12 @@
 
 package io.cdap.wrangler;
 
-import io.cdap.cdap.api.data.schema.Schema;
-import io.cdap.cdap.etl.api.relational.Capability;
-import io.cdap.cdap.etl.api.relational.Engine;
-import io.cdap.cdap.etl.api.relational.Expression;
-import io.cdap.cdap.etl.api.relational.ExpressionFactory;
-import io.cdap.cdap.etl.api.relational.ExpressionFactoryType;
 import io.cdap.cdap.etl.api.relational.InvalidRelation;
 import io.cdap.cdap.etl.api.relational.Relation;
 
-import io.cdap.cdap.etl.api.relational.RelationalTranformContext;
-import io.cdap.directives.transformation.Lower;
-import io.cdap.wrangler.api.Directive;
-import io.cdap.wrangler.api.DirectiveLoadException;
-import io.cdap.wrangler.api.DirectiveParseException;
-import io.cdap.wrangler.api.RecipeException;
-import io.cdap.wrangler.parser.GrammarBasedParser;
-import io.cdap.wrangler.parser.MigrateToV2;
-import io.cdap.wrangler.registry.DirectiveRegistry;
-import io.cdap.wrangler.registry.SystemDirectiveRegistry;
-import io.cdap.wrangler.utils.SqlExpressionGenerator;
-import org.apache.poi.ss.formula.functions.T;
 import org.junit.Assert;
 import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.InjectMocks;
-import org.mockito.Mock;
-import org.mockito.Mockito;
-import org.mockito.junit.MockitoJUnitRunner;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collection;
-import java.util.List;
-import java.util.Optional;
-import java.util.Set;
-
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;