Add tests for data source option for pivot tables

2 files changed, 18 insertions, 0 deletions

diff --git a/test/workbook/worksheet/tc_pivot_table.rb b/test/workbook/worksheet/tc_pivot_table.rb
index ee90bec2..56ac4bd4 100644
--- a/test/workbook/worksheet/tc_pivot_table.rb
+++ b/test/workbook/worksheet/tc_pivot_table.rb
@@ -36,6 +36,16 @@ class TestPivotTable < Test::Unit::TestCase
     assert_equal(@ws.pivot_tables.last, pivot_table, "must be added to worksheet pivot tables collection")
   end
 
+  def test_set_pivot_table_data_sheet
+    pivot_table = @ws.add_pivot_table('G5:G6', 'A1:D5')
+    data_sheet = @ws.clone
+    data_sheet.name = "Pivot Table Data Source"
+
+    assert_equal(pivot_table.data_sheet.name, @ws.name, "must default to the same sheet the pivot table is added to")
+    pivot_table.data_sheet = data_sheet
+    assert_equal(pivot_table.data_sheet.name, data_sheet.name, "data sheet assigned to pivot table")
+  end
+
   def test_add_pivot_table_with_config
     pivot_table = @ws.add_pivot_table('G5:G6', 'A1:D5') do |pt|
       pt.rows = ['Year', 'Month']
diff --git a/test/workbook/worksheet/tc_pivot_table_cache_definition.rb b/test/workbook/worksheet/tc_pivot_table_cache_definition.rb
index a38e808a..282078b0 100644
--- a/test/workbook/worksheet/tc_pivot_table_cache_definition.rb
+++ b/test/workbook/worksheet/tc_pivot_table_cache_definition.rb
@@ -32,6 +32,14 @@ class TestPivotTableCacheDefinition < Test::Unit::TestCase
     assert_equal(1, @cache_definition.cache_id)
   end
 
+  def test_data_sheet
+    data_sheet = @ws.clone
+    data_sheet.name = "Pivot Table Data Source"
+    @pivot_table.data_sheet = data_sheet
+
+    assert(@cache_definition.to_xml_string.include?(data_sheet.name), "must set the data source correctly")
+  end
+
   def test_to_xml_string
     schema = Nokogiri::XML::Schema(File.open(Axlsx::SML_XSD))
     doc = Nokogiri::XML(@cache_definition.to_xml_string)