PY-20932 Pandas DataFrame viewer for multiindex frames raises TypeError
[idea/community.git] / 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 }