From 3aa573bdd866a8f2ab814533a4d3623d2d4bd5ff Mon Sep 17 00:00:00 2001
From: PJ Fanning <fanningpj@apache.org>
Date: Fri, 3 Jul 2020 17:02:59 +0000
Subject: [PATCH] add test case

git-svn-id: https://svn.apache.org/repos/asf/poi/trunk@1879485 13f79535-47bb-0310-9956-ffa450edef68
---
 .../xssf/streaming/DeferredSXSSFWorkbook.java | 36 ++-----------------
 .../poi/xssf/streaming/SXSSFWorkbook.java     | 19 ++++------
 .../streaming/TestDeferredSXSSFWorkbook.java  | 28 +++++++++++++--
 3 files changed, 35 insertions(+), 48 deletions(-)

diff --git a/src/ooxml/java/org/apache/poi/xssf/streaming/DeferredSXSSFWorkbook.java b/src/ooxml/java/org/apache/poi/xssf/streaming/DeferredSXSSFWorkbook.java
index 6407adf73c..780ebf0c80 100644
--- a/src/ooxml/java/org/apache/poi/xssf/streaming/DeferredSXSSFWorkbook.java
+++ b/src/ooxml/java/org/apache/poi/xssf/streaming/DeferredSXSSFWorkbook.java
@@ -101,41 +101,9 @@ public class DeferredSXSSFWorkbook extends SXSSFWorkbook {
      */
     @Override
     public Iterator<Sheet> sheetIterator() {
-        return new SheetIterator<Sheet>();
+        return new SheetIterator<>();
     }
-    
-    private final class SheetIterator<T extends Sheet> implements Iterator<T> {
-        final private Iterator<XSSFSheet> it;
-        
-        @SuppressWarnings("unchecked")
-        public SheetIterator() {
-            it = (Iterator<XSSFSheet>) (Iterator<? extends Sheet>) _wb.iterator();
-        }
-        
-        @Override
-        public boolean hasNext() {
-            return it.hasNext();
-        }
-        
-        @Override
-        @SuppressWarnings("unchecked")
-        public T next() throws NoSuchElementException {
-            final XSSFSheet xssfSheet = it.next();
-            DeferredSXSSFSheet sxSheet = (DeferredSXSSFSheet) getSXSSFSheet(xssfSheet);
-            return (T) (sxSheet == null ? xssfSheet : sxSheet);
-        }
-        
-        /**
-         * Unexpected behavior may occur if sheets are reordered after iterator has been created. Support for the remove
-         * method may be added in the future if someone can figure out a reliable implementation.
-         */
-        @Override
-        public void remove() throws IllegalStateException {
-            throw new UnsupportedOperationException("remove method not supported on XSSFWorkbook.iterator(). "
-                    + "Use Sheet.removeSheetAt(int) instead.");
-        }
-    }
-    
+
     /**
      * Alias for {@link #sheetIterator()} to allow foreach loops
      */
diff --git a/src/ooxml/java/org/apache/poi/xssf/streaming/SXSSFWorkbook.java b/src/ooxml/java/org/apache/poi/xssf/streaming/SXSSFWorkbook.java
index 4c714a4adb..8e6445ffeb 100644
--- a/src/ooxml/java/org/apache/poi/xssf/streaming/SXSSFWorkbook.java
+++ b/src/ooxml/java/org/apache/poi/xssf/streaming/SXSSFWorkbook.java
@@ -367,18 +367,13 @@ public class SXSSFWorkbook implements Workbook {
 
     void deregisterSheetMapping(XSSFSheet xSheet)
     {
-        SXSSFSheet sxSheet=getSXSSFSheet(xSheet);
-
-        // ensure that the writer is closed in all cases to not have lingering writers
-        try {
-            sxSheet.getSheetDataWriter().close();
-        } catch (IOException e) {
-            // ignore exception here
+        SXSSFSheet sxSheet = getSXSSFSheet(xSheet);
+        if (sxSheet != null) {
+            // ensure that the writer is closed in all cases to not have lingering writers
+            IOUtils.closeQuietly(sxSheet.getSheetDataWriter());
+            _sxFromXHash.remove(sxSheet);
+            _xFromSxHash.remove(xSheet);
         }
-
-        _sxFromXHash.remove(sxSheet);
-
-        _xFromSxHash.remove(xSheet);
     }
 
     protected XSSFSheet getSheetFromZipEntryName(String sheetRef)
@@ -755,7 +750,7 @@ public class SXSSFWorkbook implements Workbook {
         return new SheetIterator<>();
     }
 
-    private final class SheetIterator<T extends Sheet> implements Iterator<T> {
+    protected final class SheetIterator<T extends Sheet> implements Iterator<T> {
         final private Iterator<XSSFSheet> it;
         @SuppressWarnings("unchecked")
         public SheetIterator() {
diff --git a/src/ooxml/testcases/org/apache/poi/xssf/streaming/TestDeferredSXSSFWorkbook.java b/src/ooxml/testcases/org/apache/poi/xssf/streaming/TestDeferredSXSSFWorkbook.java
index 2d3026fdd1..ffdbc0516f 100644
--- a/src/ooxml/testcases/org/apache/poi/xssf/streaming/TestDeferredSXSSFWorkbook.java
+++ b/src/ooxml/testcases/org/apache/poi/xssf/streaming/TestDeferredSXSSFWorkbook.java
@@ -199,12 +199,36 @@ public final class TestDeferredSXSSFWorkbook extends BaseTestXWorkbook {
         assertNull(wr);
         wb.close();
     }
+
+    @Test
+    public void removeSheet() throws IOException {
+        try (DeferredSXSSFWorkbook wb = new DeferredSXSSFWorkbook()) {
+            DeferredSXSSFSheet sheet1 = wb.createSheet("sheet1");
+            sheet1.setRowGenerator((sh) -> {
+                Row row = sh.createRow(0);
+                Cell cell = row.createCell(0);
+                cell.setCellValue("sheet1");
+            });
+            DeferredSXSSFSheet sheet2 = wb.createSheet("sheet2");
+            sheet2.setRowGenerator((sh) -> {
+                Row row = sh.createRow(0);
+                Cell cell = row.createCell(0);
+                cell.setCellValue("sheet2");
+            });
+            wb.removeSheetAt(0);
+            try (XSSFWorkbook wb2 = SXSSFITestDataProvider.instance.writeOutAndReadBack(wb)) {
+                assertNull(wb2.getSheet(  "sheet1"));
+                XSSFSheet xssfSheet = wb2.getSheet(  "sheet2");
+                assertNotNull(xssfSheet);
+                assertEquals("sheet2", xssfSheet.getRow(0).getCell(0).getStringCellValue());
+            }
+        }
+    }
     
     @Test
     public void gzipSheetdataWriter() throws IOException {
         DeferredSXSSFWorkbook wb = new DeferredSXSSFWorkbook();
-        wb.setCompressTempFiles(true);
-        
+
         final int rowNum = 1000;
         final int sheetNum = 5;
         populateData(wb, 1000, 5);
-- 
2.39.5