diff --git a/package.json b/package.json
index 1916e7c..6e70bcf 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
     "name": "dash-ag-grid",
-    "version": "31.2.1",
+    "version": "31.0.1",
     "description": "Dash wrapper around AG Grid, the best interactive data grid for the web.",
     "repository": {
         "type": "git",
diff --git a/tests/test_selected_rows.py b/tests/test_selected_rows.py
index b156bd9..53fdafd 100644
--- a/tests/test_selected_rows.py
+++ b/tests/test_selected_rows.py
@@ -110,6 +110,7 @@ def test_sr2_selected_rows_rowdata(dash_duo):
         dag.AgGrid(
             id='grid',
             columnSize="sizeToFit",
+            columnDefs=[{'field': 'A'}],
             dashGridOptions={
                 "rowHeight": None,
                 "domLayout": "normal",
@@ -152,6 +153,8 @@ def update_selected_row_info(selected_rows):
 
     grid = utils.Grid(dash_duo, "grid")
 
+    grid.wait_for_header_text(0, 'A')
+
     dash_duo.find_element('#update-button').click()
 
     grid.wait_for_cell_text(0, 0, "1")
@@ -186,6 +189,7 @@ def test_sr3_selected_rows_modes(dash_duo):
         dag.AgGrid(
             id='grid',
             columnSize="sizeToFit",
+            columnDefs=[{'field': 'A'}],
             dashGridOptions={
                 "rowHeight": None,
                 "domLayout": "normal",
@@ -245,6 +249,8 @@ def update_selected_row_info(selected_rows):
 
     grid = utils.Grid(dash_duo, "grid")
 
+    grid.wait_for_header_text(0, 'A')
+
     dash_duo.find_element('#update-button').click()
 
     grid.wait_for_cell_text(0, 0, "1")