From a1e1db1e6f2faa8286afb1e9291c51fa084b66f7 Mon Sep 17 00:00:00 2001
From: xuri <xuri.me@gmail.com>
Date: Fri, 14 May 2021 00:09:23 +0800
Subject: [PATCH] This closes #838, fix wrong worksheet XML path of the stream
 writer in some case

---
 stream.go      | 16 ++++++++++------
 stream_test.go |  6 +++---
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/stream.go b/stream.go
index f12b201f..e5fa2373 100644
--- a/stream.go
+++ b/stream.go
@@ -99,11 +99,11 @@ func (f *File) NewStreamWriter(sheet string) (*StreamWriter, error) {
 		return nil, err
 	}
 
-	sheetXML := fmt.Sprintf("xl/worksheets/sheet%d.xml", sw.SheetID)
+	sheetPath := f.sheetMap[trimSheetName(sheet)]
 	if f.streams == nil {
 		f.streams = make(map[string]*StreamWriter)
 	}
-	f.streams[sheetXML] = sw
+	f.streams[sheetPath] = sw
 
 	_, _ = sw.rawData.WriteString(XMLHeader + `<worksheet` + templateNamespaceIDMap)
 	bulkAppendFields(&sw.rawData, sw.worksheet, 2, 5)
@@ -462,6 +462,10 @@ func writeCell(buf *bufferedWriter, c xlsxC) {
 
 // Flush ending the streaming writing process.
 func (sw *StreamWriter) Flush() error {
+	if !sw.sheetWritten {
+		_, _ = sw.rawData.WriteString(`<sheetData>`)
+		sw.sheetWritten = true
+	}
 	_, _ = sw.rawData.WriteString(`</sheetData>`)
 	bulkAppendFields(&sw.rawData, sw.worksheet, 8, 15)
 	if sw.mergeCellsCount > 0 {
@@ -476,10 +480,10 @@ func (sw *StreamWriter) Flush() error {
 		return err
 	}
 
-	sheetXML := fmt.Sprintf("xl/worksheets/sheet%d.xml", sw.SheetID)
-	delete(sw.File.Sheet, sheetXML)
-	delete(sw.File.checked, sheetXML)
-	delete(sw.File.XLSX, sheetXML)
+	sheetPath := sw.File.sheetMap[trimSheetName(sw.Sheet)]
+	delete(sw.File.Sheet, sheetPath)
+	delete(sw.File.checked, sheetPath)
+	delete(sw.File.XLSX, sheetPath)
 
 	return nil
 }
diff --git a/stream_test.go b/stream_test.go
index 0834a2dc..391c99d0 100644
--- a/stream_test.go
+++ b/stream_test.go
@@ -68,13 +68,13 @@ func TestStreamWriter(t *testing.T) {
 		assert.NoError(t, streamWriter.SetRow(cell, row))
 	}
 
-	// Test set cell column overflow.
-	assert.EqualError(t, streamWriter.SetRow("XFD1", []interface{}{"A", "B", "C"}), ErrColumnNumber.Error())
-
 	assert.NoError(t, streamWriter.Flush())
 	// Save spreadsheet by the given path.
 	assert.NoError(t, file.SaveAs(filepath.Join("test", "TestStreamWriter.xlsx")))
 
+	// Test set cell column overflow.
+	assert.EqualError(t, streamWriter.SetRow("XFD1", []interface{}{"A", "B", "C"}), ErrColumnNumber.Error())
+
 	// Test close temporary file error.
 	file = NewFile()
 	streamWriter, err = file.NewStreamWriter("Sheet1")