python/testSrc/com/jetbrains/env/python/PythonDataViewerTest.java

   1 /*
   2  * Copyright 2000-2016 JetBrains s.r.o.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  * http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 package com.jetbrains.env.python;
  17
  18 import com.google.common.collect.ImmutableSet;
  19 import com.intellij.util.Consumer;
  20 import com.intellij.xdebugger.XDebugSession;
  21 import com.intellij.xdebugger.XDebuggerTestUtil;
  22 import com.jetbrains.env.PyEnvTestCase;
  23 import com.jetbrains.env.Staging;
  24 import com.jetbrains.env.python.debug.PyDebuggerTask;
  25 import com.jetbrains.python.debugger.ArrayChunk;
  26 import com.jetbrains.python.debugger.PyDebugValue;
  27 import com.jetbrains.python.debugger.PyDebuggerException;
  28 import org.jetbrains.annotations.NotNull;
  29 import org.jetbrains.annotations.Nullable;
  30 import org.junit.Test;
  31
  32 import java.lang.reflect.InvocationTargetException;
  33 import java.util.List;
  34 import java.util.Set;
  35
  36 import static com.intellij.testFramework.UsefulTestCase.assertSameElements;
  37 import static org.junit.Assert.assertEquals;
  38
  39 /**
  40  * Created by Yuli Fiterman on 5/10/2016.
  41  */
  42 public class PythonDataViewerTest extends PyEnvTestCase {
  43
  44   @Test
  45   @Staging
  46   public void testDataFrameChunkRetrieval() throws Exception {
  47     runPythonTest(new PyDataFrameDebuggerTask(getRelativeTestDataPath(), "test_dataframe.py", ImmutableSet.of(7, 15, 22)) {
  48       @Override
  49       public void testing() throws Exception {
  50         doTest("df1", 3, 5, null);
  51
  52         doTest("df2", 3, 6, arrayChunk -> {
  53           List<ArrayChunk.ColHeader> colHeaders = arrayChunk.getColHeaders();
  54           assertSameElements(colHeaders.stream().map(ArrayChunk.ColHeader::getLabel).toArray(),
  55                              "LABELS", "One_X", "One_Y", "Two_X", "Two_Y", "row");
  56         });
  57
  58         doTest("df3", 7, 3, arrayChunk -> {
  59           ArrayChunk.ColHeader header = arrayChunk.getColHeaders().get(2);
  60           assertEquals("Sales", header.getLabel());
  61           assertEquals(16, (int)Integer.valueOf(header.getMax()));
  62           assertEquals(1, (int)Integer.valueOf(header.getMin()));
  63         });
  64       }
  65     });
  66   }
  67
  68   @Test
  69   @Staging
  70   public void testMultiIndexDataFrame() throws Exception {
  71     runPythonTest(new PyDataFrameDebuggerTask(getRelativeTestDataPath(), "test_dataframe_multiindex.py", ImmutableSet.of(5, 10)) {
  72       @Override
  73       public void testing() throws Exception {
  74         doTest("frame1", 4, 2, arrayChunk -> assertSameElements(arrayChunk.getRowLabels(),
  75                                                                 "s/2", "s/3", "d/2", "d/3"));
  76         doTest("frame2", 4, 4, arrayChunk -> {
  77           List<ArrayChunk.ColHeader> headers = arrayChunk.getColHeaders();
  78           assertSameElements(headers.stream().map(ArrayChunk.ColHeader::getLabel).toArray(), "1/1", "1/B", "2/1", "2/B");
  79         });
  80       }
  81     });
  82   }
  83
  84   private static class PyDataFrameDebuggerTask extends PyDebuggerTask {
  85
  86     private Set<Integer> myLines;
  87
  88     public PyDataFrameDebuggerTask(@Nullable String relativeTestDataPath, String scriptName, Set<Integer> lines) {
  89       super(relativeTestDataPath, scriptName);
  90       myLines = lines;
  91     }
  92
  93     protected void testShape(ArrayChunk arrayChunk, int expectedRows, int expectedColumns) {
  94       assertEquals(expectedRows, arrayChunk.getRows());
  95       assertEquals(expectedColumns, arrayChunk.getColumns());
  96     }
  97
  98     protected void doTest(String name, int expectedRows, int expectedColumns, @Nullable Consumer<ArrayChunk> test)
  99       throws InvocationTargetException, InterruptedException, PyDebuggerException {
 100       waitForPause();
 101       ArrayChunk arrayChunk = getDefaultChunk(name, mySession);
 102       testShape(arrayChunk, expectedRows, expectedColumns);
 103       if (test != null) {
 104         test.consume(arrayChunk);
 105       }
 106       resume();
 107     }
 108
 109     @Override
 110     public void before() throws Exception {
 111       for (Integer line : myLines) {
 112         toggleBreakpoint(getScriptName(), line);
 113       }
 114     }
 115
 116     @NotNull
 117     @Override
 118     public Set<String> getTags() {
 119       return ImmutableSet.of("pandas");
 120     }
 121   }
 122
 123   private static ArrayChunk getDefaultChunk(String varName, XDebugSession session) throws PyDebuggerException {
 124     PyDebugValue dbgVal = (PyDebugValue)XDebuggerTestUtil.evaluate(session, varName).first;
 125     return dbgVal.getFrameAccessor().getArrayItems(dbgVal, 0, 0, -1, -1, ".%5f");
 126   }
 127
 128   private static String getRelativeTestDataPath() {
 129     return "/debug";
 130   }
 131 }