d988bb361686a6c0aa1e8116aa71e9d9c1568f03
[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.xdebugger.XDebuggerTestUtil;
20 import com.jetbrains.env.PyEnvTestCase;
21 import com.jetbrains.env.Staging;
22 import com.jetbrains.env.python.debug.PyDebuggerTask;
23 import com.jetbrains.python.debugger.ArrayChunk;
24 import com.jetbrains.python.debugger.PyDebugValue;
25 import com.jetbrains.python.debugger.PyDebuggerException;
26 import com.jetbrains.python.debugger.array.AsyncArrayTableModel;
27 import org.jetbrains.annotations.NotNull;
28 import org.junit.Test;
29
30 import java.util.Set;
31
32 import static com.intellij.testFramework.UsefulTestCase.assertSameElements;
33 import static org.junit.Assert.assertEquals;
34
35 /**
36  * Created by Yuli Fiterman on 5/10/2016.
37  */
38 public class PythonDataViewerTest extends PyEnvTestCase {
39
40   @Test
41   @Staging
42   public void testDataFrameChunkRetrieval() throws Exception {
43     runPythonTest(new PyDebuggerTask("/debug", "test_dataframe.py") {
44       @Override
45       public void before() throws Exception {
46         toggleBreakpoint(getScriptName(), 7);
47         toggleBreakpoint(getScriptName(), 15);
48         toggleBreakpoint(getScriptName(), 22);
49       }
50
51       @Override
52       public void testing() throws Exception {
53         waitForPause();
54         ArrayChunk df1 = getDefaultChunk("df1");
55         assertEquals(5, df1.getColumns());
56         assertEquals(3, df1.getRows());
57         resume();
58
59         waitForPause();
60         ArrayChunk df2 = getDefaultChunk("df2");
61         assertEquals(6, df2.getColumns());
62         assertEquals(3, df2.getRows());
63         assertSameElements(df2.getColHeaders().stream().map((header -> header.getLabel())).toArray(),
64                            new String[]{"LABELS", "One_X", "One_Y", "Two_X", "Two_Y", "row"});
65         resume();
66
67         waitForPause();
68         ArrayChunk df3 = getDefaultChunk("df3");
69         assertEquals(3, df3.getColumns());
70         assertEquals(7, df3.getRows());
71         ArrayChunk.ColHeader header = df3.getColHeaders().get(2);
72         assertEquals("Sales", header.getLabel());
73         assertEquals(16, (int)Integer.valueOf(header.getMax()));
74         assertEquals(1, (int)Integer.valueOf(header.getMin()));
75         resume();
76       }
77
78       private ArrayChunk getDefaultChunk(String varName) throws PyDebuggerException {
79         PyDebugValue dbgVal = (PyDebugValue)XDebuggerTestUtil.evaluate(mySession, varName).first;
80         return dbgVal.getFrameAccessor()
81           .getArrayItems(dbgVal, 0, 0, -1, -1, ".%5f");
82       }
83
84
85       @NotNull
86       @Override
87       public Set<String> getTags() {
88         return ImmutableSet.of("pandas");
89       }
90     });
91   }
92 }