Convert DataSize for tests

Author: shangm2

presto-main/src/test/java/com/facebook/presto/execution/MockManagedQueryExecution.java

@@ -42,15 +42,14 @@
 import static com.facebook.presto.execution.QueryState.WAITING_FOR_PREREQUISITES;
 import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
 import static io.airlift.units.DataSize.Unit.BYTE;
-import static io.airlift.units.DataSize.succinctBytes;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 import static java.util.concurrent.TimeUnit.NANOSECONDS;
 
 public class MockManagedQueryExecution
         implements ManagedQueryExecution
 {
     private final List<StateChangeListener<QueryState>> listeners = new ArrayList<>();
-    private final DataSize memoryUsage;
+    private final long memoryUsage;
     private final Duration cpuUsage;
     private final Session session;
     private QueryState state = WAITING_FOR_PREREQUISITES;
@@ -75,7 +74,7 @@ public MockManagedQueryExecution(long memoryUsage, String queryId, int priority,
 
     public MockManagedQueryExecution(long memoryUsage, String queryId, int priority, Duration cpuUsage, ResourceGroupId resourceGroupId)
     {
-        this.memoryUsage = succinctBytes(memoryUsage);
+        this.memoryUsage = memoryUsage;
         this.cpuUsage = cpuUsage;
         this.session = testSessionBuilder()
                 .setSystemProperty(QUERY_PRIORITY, String.valueOf(priority))
@@ -161,13 +160,13 @@ public BasicQueryInfo getBasicQueryInfo()
     }
 
     @Override
-    public DataSize getUserMemoryReservation()
+    public long getUserMemoryReservationInBytes()
     {
         return memoryUsage;
     }
 
     @Override
-    public DataSize getTotalMemoryReservation()
+    public long getTotalMemoryReservationInBytes()
     {
         return memoryUsage;
     }

presto-main/src/test/java/com/facebook/presto/execution/MockQueryExecution.java

@@ -20,7 +20,6 @@
 import com.facebook.presto.spi.resourceGroups.ResourceGroupQueryLimits;
 import com.facebook.presto.sql.planner.Plan;
 import com.google.common.util.concurrent.ListenableFuture;
-import io.airlift.units.DataSize;
 import io.airlift.units.Duration;
 import org.joda.time.DateTime;
 
@@ -67,11 +66,11 @@ public ListenableFuture<QueryState> getStateChange(QueryState currentState)
 
     @Override
     public void addStateChangeListener(StateMachine.StateChangeListener<QueryState> stateChangeListener)
-    { }
+    {}
 
     @Override
     public void addOutputInfoListener(Consumer<QueryOutputInfo> listener)
-    { }
+    {}
 
     @Override
     public Plan getQueryPlan()
@@ -110,15 +109,15 @@ public Duration getTotalCpuTime()
     }
 
     @Override
-    public DataSize getRawInputDataSize()
+    public long getRawInputDataSizeInBytes()
     {
-        return null;
+        return 0;
     }
 
     @Override
-    public DataSize getWrittenIntermediateDataSize()
+    public long getWrittenIntermediateDataSizeInBytes()
     {
-        return null;
+        return 0;
     }
 
     @Override
@@ -128,9 +127,9 @@ public long getOutputPositions()
     }
 
     @Override
-    public DataSize getOutputDataSize()
+    public long getOutputDataSizeInBytes()
     {
-        return null;
+        return 0;
     }
 
     @Override
@@ -140,15 +139,15 @@ public int getRunningTaskCount()
     }
 
     @Override
-    public DataSize getUserMemoryReservation()
+    public long getUserMemoryReservationInBytes()
     {
-        return null;
+        return 0;
     }
 
     @Override
-    public DataSize getTotalMemoryReservation()
+    public long getTotalMemoryReservationInBytes()
     {
-        return null;
+        return 0;
     }
 
     @Override
@@ -218,37 +217,37 @@ public Optional<Throwable> getFailureReason()
 
     @Override
     public void pruneExpiredQueryInfo()
-    { }
+    {}
 
     @Override
     public void pruneFinishedQueryInfo()
-    { }
+    {}
 
     @Override
     public void setResourceGroupQueryLimits(ResourceGroupQueryLimits resourceGroupQueryLimits)
-    { }
+    {}
 
     @Override
     public void setMemoryPool(VersionedMemoryPoolId poolId)
-    { }
+    {}
 
     @Override
     public void start()
-    { }
+    {}
 
     @Override
     public void cancelQuery()
-    { }
+    {}
 
     @Override
     public void cancelStage(StageId stageId)
-    { }
+    {}
 
     @Override
     public void recordHeartbeat()
-    { }
+    {}
 
     @Override
     public void addFinalQueryInfoListener(StateMachine.StateChangeListener<QueryInfo> stateChangeListener)
-    { }
+    {}
 }

presto-main/src/test/java/com/facebook/presto/execution/MockRemoteTaskFactory.java

@@ -92,7 +92,6 @@
 import static com.facebook.presto.util.Failures.toFailures;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.util.concurrent.Futures.nonCancellationPropagating;
-import static io.airlift.units.DataSize.Unit.BYTE;
 import static io.airlift.units.DataSize.Unit.GIGABYTE;
 import static io.airlift.units.DataSize.Unit.MEGABYTE;
 import static java.lang.Math.addExact;
@@ -240,7 +239,7 @@ public MockRemoteTask(TaskId taskId,
                     taskId,
                     TASK_INSTANCE_ID.toString(),
                     executor,
-                    new DataSize(1, BYTE),
+                    1L,
                     () -> new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), "test"),
                     new SpoolingOutputBufferFactory(new FeaturesConfig()));
 
@@ -444,9 +443,9 @@ public synchronized void noMoreSplits(PlanNodeId sourceId)
             noMoreSplits.add(sourceId);
 
             boolean allSourcesComplete = Stream.concat(
-                    fragment.getTableScanSchedulingOrder().stream(),
-                    fragment.getRemoteSourceNodes().stream()
-                            .map(PlanNode::getId))
+                            fragment.getTableScanSchedulingOrder().stream(),
+                            fragment.getRemoteSourceNodes().stream()
+                                    .map(PlanNode::getId))
                     .allMatch(noMoreSplits::contains);
 
             if (allSourcesComplete) {

presto-main/src/test/java/com/facebook/presto/execution/TestMemoryRevokingScheduler.java

@@ -796,7 +796,7 @@ private SqlTask newSqlTask(QueryId queryId, MemoryPool memoryPool)
                 new MockExchangeClientSupplier(),
                 singleThreadedExecutor,
                 Functions.identity(),
-                new DataSize(32, MEGABYTE),
+                new DataSize(32, MEGABYTE).toBytes(),
                 new CounterStat(),
                 new SpoolingOutputBufferFactory(new FeaturesConfig()));
     }

presto-main/src/test/java/com/facebook/presto/execution/TestQueryLimit.java

@@ -33,13 +33,13 @@
 public class TestQueryLimit
 {
     private static final QueryLimit<Duration> QUERY_LIMIT_DURATION = createDurationLimit(new Duration(1, HOURS), SYSTEM);
-    private static final QueryLimit<DataSize> QUERY_LIMIT_DATA_SIZE = createDataSizeLimit(new DataSize(1, MEGABYTE), RESOURCE_GROUP);
+    private static final QueryLimit<Long> QUERY_LIMIT_DATA_SIZE = createDataSizeLimit(new DataSize(1, MEGABYTE).toBytes(), RESOURCE_GROUP);
 
     @Test
     public void testGetLimit()
     {
         assertEquals(QUERY_LIMIT_DURATION.getLimit(), new Duration(1, HOURS));
-        assertEquals(QUERY_LIMIT_DATA_SIZE.getLimit(), new DataSize(1, MEGABYTE));
+        assertEquals(QUERY_LIMIT_DATA_SIZE.getLimit(), new DataSize(1, MEGABYTE).toBytes());
     }
 
     @Test
@@ -57,8 +57,8 @@ public void testGetMinimum()
         QueryLimit<Duration> shortLimit = createDurationLimit(new Duration(1, MINUTES), RESOURCE_GROUP);
         assertEquals(getMinimum(QUERY_LIMIT_DURATION, longLimit, shortLimit), shortLimit);
 
-        QueryLimit<DataSize> largeLimit = createDataSizeLimit(new DataSize(2, MEGABYTE), QUERY);
-        QueryLimit<DataSize> smallLimit = createDataSizeLimit(new DataSize(1, BYTE), RESOURCE_GROUP);
+        QueryLimit<Long> largeLimit = createDataSizeLimit(new DataSize(2, MEGABYTE).toBytes(), QUERY);
+        QueryLimit<Long> smallLimit = createDataSizeLimit(new DataSize(1, BYTE).toBytes(), RESOURCE_GROUP);
         assertEquals(getMinimum(smallLimit, QUERY_LIMIT_DATA_SIZE, largeLimit), smallLimit);
         assertEquals(getMinimum(null, QUERY_LIMIT_DATA_SIZE, largeLimit), QUERY_LIMIT_DATA_SIZE);
         assertEquals(getMinimum(null, largeLimit, null, null, QUERY_LIMIT_DATA_SIZE), QUERY_LIMIT_DATA_SIZE);