--- a/usr/src/cmd/text-install/osol_install/text_install/disk_selection.py	Fri May 14 13:30:56 2010 -0700
+++ b/usr/src/cmd/text-install/osol_install/text_install/disk_selection.py	Fri May 21 08:53:04 2010 -0600
@@ -18,8 +18,7 @@
 #
 # CDDL HEADER END
 #
-# Copyright 2010 Sun Microsystems, Inc.  All rights reserved.
-# Use is subject to license terms.
+# Copyright (c) 2009, 2010, Oracle and/or its affiliates. All rights reserved.
 #
 
 '''
@@ -98,11 +97,6 @@
     
     def __init__(self, main_win):
         super(DiskScreen, self).__init__(main_win)
-        self.recommended_size = get_recommended_size().size_as("gb")
-        self.minimum_size = get_minimum_size().size_as("gb")
-        size_dict = {"recommend" : self.recommended_size,
-                     "min" : self.minimum_size}
-        self.size_line = DiskScreen.SIZE_TEXT % size_dict
         if platform.processor() == "i386":
             self.found_text = DiskScreen.FOUND_x86
             self.proposed_text = DiskScreen.PROPOSED_x86
@@ -130,6 +124,45 @@
         self.disk_detail = None
         self.num_targets = 0
         self.td_handle = None
+        self._size_line = None
+        self.selected_disk = 0
+        self._minimum_size = None
+        self._recommended_size = None
+        self.do_copy = False # Flag indicating if install_profile.disk
+                             # should be copied
+    
+    def determine_minimum(self):
+        '''Returns minimum install size, fetching first if needed'''
+        self.determine_size_data()
+        return self._minimum_size
+    
+    minimum_size = property(determine_minimum)
+    
+    def determine_recommended(self):
+        '''Returns recommended install size, fetching first if needed'''
+        self.determine_size_data()
+        return self._recommended_size
+    
+    recommended_size = property(determine_recommended)
+    
+    def determine_size_data(self):
+        '''Retrieve the minimum and recommended sizes and generate the string
+        to present that information.
+        
+        '''
+        if self._minimum_size is None or self._recommended_size is None:
+            self._recommended_size = get_recommended_size().size_as("gb")
+            self._minimum_size = get_minimum_size().size_as("gb")
+    
+    def get_size_line(self):
+        '''Returns the line of text displaying the min/recommended sizes'''
+        if self._size_line is None:
+            size_dict = {"recommend" : self.recommended_size,
+                         "min" : self.minimum_size}
+            self._size_line = DiskScreen.SIZE_TEXT % size_dict
+        return self._size_line
+    
+    size_line = property(get_size_line)
     
     def wait_for_disks(self):
         '''Block while waiting for libtd to finish. Catch F9 and quit
@@ -294,21 +327,23 @@
         
         self.main_win.do_update()
         self.center_win.activate_object(self.disk_win)
-        try:
-            self.disk_win.activate_object(self.install_profile.disk.TUI_INDEX)
-        except AttributeError:
-            self.disk_win.activate_object()
+        self.disk_win.activate_object(self.selected_disk)
+        # Set the flag so that the disk is not copied by on_change_screen,
+        # unless on_activate gets called as a result of the user changing
+        # the selected disk.
+        self.do_copy = False
     
     def on_change_screen(self):
         ''' Assign the selected disk to the InstallProfile, and make note of
         its index (in case the user returns to this screen later)
         
         '''
-        if self.num_targets > 0:
-            disk = self.disk_detail.disk_info
-            self.install_profile.disk = deepcopy(disk)
-            self.install_profile.original_disk = disk
-            self.install_profile.disk.TUI_INDEX = self.disk_win.active_object
+        if self.disk_detail is not None:
+            if self.do_copy or self.install_profile.disk is None:
+                disk = self.disk_detail.disk_info
+                self.install_profile.disk = deepcopy(disk)
+                self.install_profile.original_disk = disk
+            self.selected_disk = self.disk_win.active_object
     
     def start_discovery(self):
         '''Spawn a thread to begin target discovery'''
@@ -386,3 +421,5 @@
         disk_select.center_win.add_paragraph(disk_select.found_text, 11, 1,
                                              max_x=max_x)
     disk_select.disk_detail.set_disk_info(disk_info)
+    # User selected a different disk; set the flag so that it gets copied later
+    disk_select.do_copy = True

--- a/usr/src/cmd/text-install/osol_install/text_install/disk_window.py	Fri May 14 13:30:56 2010 -0700
+++ b/usr/src/cmd/text-install/osol_install/text_install/disk_window.py	Fri May 21 08:53:04 2010 -0600
@@ -18,8 +18,7 @@
 #
 # CDDL HEADER END
 #
-# Copyright 2010 Sun Microsystems, Inc.  All rights reserved.
-# Use is subject to license terms.
+# Copyright (c) 2009, 2010, Oracle and/or its affiliates. All rights reserved.
 #
 
 
@@ -122,7 +121,11 @@
         self.key_dict[curses.KEY_RIGHT] = self.on_arrow_key
         if self.editable:
             self.key_dict[curses.KEY_F5] = self.change_type
-        self.set_disk_info(disk_info)
+        
+        if getattr(disk_info, "do_revert", False):
+            self.reset()
+        else:
+            self.set_disk_info(disk_info)
     
     def _init_win(self, window):
         '''Require at least 70 columns and 6 lines to fit current needs for

--- a/usr/src/cmd/text-install/osol_install/text_install/fdisk_partitions.py	Fri May 14 13:30:56 2010 -0700
+++ b/usr/src/cmd/text-install/osol_install/text_install/fdisk_partitions.py	Fri May 21 08:53:04 2010 -0600
@@ -18,8 +18,7 @@
 #
 # CDDL HEADER END
 #
-# Copyright 2010 Sun Microsystems, Inc.  All rights reserved.
-# Use is subject to license terms.
+# Copyright (c) 2009, 2010, Oracle and/or its affiliates. All rights reserved.
 #
 
 '''
@@ -128,10 +127,12 @@
                 # partition
                 self.disk_info.create_default_layout()
                 raise SkipException
+            disp_disk = self.install_profile.original_disk.get_solaris_data()
             logging.debug("Preserved partition with existing slices:"
                           " presenting option to install into a slice")
         else:
             self.disk_info = self.install_profile.disk
+            disp_disk = self.install_profile.original_disk
             if self.disk_info.boot:
                 bootable = FDiskPart.BOOT_TEXT
             else:
@@ -158,7 +159,7 @@
         
         y_loc += 1
         disk_win_area = WindowArea(6, 70, y_loc, 0)
-        self.disk_win = DiskWindow(disk_win_area, self.disk_info,
+        self.disk_win = DiskWindow(disk_win_area, disp_disk,
                                    window=self.center_win)
         y_loc += disk_win_area.lines
         
@@ -200,4 +201,9 @@
         else:
             logging.debug("Setting use_whole segment false for %s",
                           type(self.disk_info))
+            # If user had previously selected to use the whole disk
+            # or partition, set the do_revert flag so that the following
+            # screen will know to reset the disk (reverting the call
+            # to create_default_layout, above)
+            self.disk_info.do_revert = self.disk_info.use_whole_segment
             self.disk_info.use_whole_segment = False

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/usr/src/cmd/text-install/osol_install/text_install/test/test_disk_select.py	Fri May 21 08:53:04 2010 -0600
@@ -0,0 +1,201 @@
+#!/usr/bin/python2.6
+#
+# CDDL HEADER START
+#
+# The contents of this file are subject to the terms of the
+# Common Development and Distribution License (the "License").
+# You may not use this file except in compliance with the License.
+#
+# You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
+# or http://www.opensolaris.org/os/licensing.
+# See the License for the specific language governing permissions
+# and limitations under the License.
+#
+# When distributing Covered Code, include this CDDL HEADER in each
+# file and include the License file at usr/src/OPENSOLARIS.LICENSE.
+# If applicable, add the following below this CDDL HEADER, with the
+# fields enclosed by brackets "[]" replaced with your own identifying
+# information: Portions Copyright [yyyy] [name of copyright owner]
+#
+# CDDL HEADER END
+#
+# Copyright (c) 2010, Oracle and/or its affiliates. All rights reserved.
+#
+
+'''
+To run these tests:
+
+1) nightly -n developer.sh # build the gate
+2) export PYTHONPATH=${WS}/proto/root_i386/usr/snadm/lib:${WS}/proto/root_i386/usr/lib/python2.6/vendor-packages
+3) pfexec python2.6 test_disk_select.py
+
+A single test may be run by specifying the test as an argument to step 3, e.g.:
+pfexec python2.6 test_disk_select.py OnActivateTest.test_on_activate_default
+
+Since the proto area is used for the PYTHONPATH, the gate must be rebuilt for
+these tests to pick up any changes in the tested code.
+
+'''
+
+import numbers
+import unittest
+
+import osol_install.text_install.disk_selection as disk_selection
+from osol_install.profile.disk_info import DiskInfo
+
+class MockCenterWin(object):
+    '''Mocks an InnerWindow as used by a MainWindow'''
+    
+    def add_paragraph(self, *args, **kwargs):
+        pass
+
+class MockDiskInfo(object):
+    '''Mocks a DiskInfo object'''
+    do_copy = False
+    label = []
+    was_blank = False
+    use_whole_segment = False
+    
+    def create_default_layout(self):
+        pass
+
+class MockDiskDetail(object):
+    '''Mocks a DiskWindow object'''
+    
+    def set_disk_info(self, *args):
+        pass
+
+class MockDiskScreen(object):
+    '''Mocks the DiskScreen'''
+    win_size_x = 0
+    proposed_text = ""
+    found_text = ""
+
+class MockInstallProfile(object):
+    '''Mocks an InstallProfile'''
+    
+    disk = None
+    original_disk = None
+
+class MockAll(object):
+    '''Generic Mock object that 'never' raises an AttributeError'''
+    
+    def __getattr__(self, name):
+        return self
+    
+    def __call__(self, *args, **kwargs):
+        return None
+
+class OnActivateTest(unittest.TestCase):
+    '''Test disk_selection.on_activate'''
+    
+    def setUp(self):
+        self.screen = MockDiskScreen()
+        self.disk = MockDiskInfo()
+        self.screen.center_win = MockCenterWin()
+        self.screen.disk_detail = MockDiskDetail()
+    
+    def tearDown(self):
+        self.screen = None
+        self.disk = None
+    
+    def test_on_activate_default(self):
+        '''Ensure that do_copy flag is set after calls to on_activate'''
+        disk_selection.on_activate(disk_info=self.disk,
+                                   disk_select=self.screen)
+        self.assertFalse(self.disk.use_whole_segment)
+        self.assertTrue(self.screen.do_copy)
+    
+    def test_on_activate_GPT(self):
+        '''Ensure use_whole_segment is set if the disk was GPT labeled'''
+        
+        self.disk.label = [DiskInfo.GPT]
+        
+        disk_selection.on_activate(disk_info=self.disk,
+                                   disk_select=self.screen)
+        self.assertTrue(self.disk.use_whole_segment)
+    
+    def test_on_activate_was_blank(self):
+        '''Ensure use_whole_segment is set when the disk was initially blank'''
+        self.disk.was_blank = True
+        
+        disk_selection.on_activate(disk_info=self.disk,
+                                   disk_select=self.screen)
+        self.assertTrue(self.disk.use_whole_segment)
+
+
+class DiskSelectTest(unittest.TestCase):
+    '''Test the DiskScreen'''
+    
+    def setUp(self):
+        self.screen = disk_selection.DiskScreen(MockAll())
+        self.screen.disk_win = MockAll()
+    
+    def test_on_change_screen_disk_detail_none(self):
+        '''Ensure selected_disk is set properly by on_change_screen'''
+        self.screen.disk_detail = None
+        obj = object()
+        self.screen.selected_disk = obj
+        
+        self.screen.on_change_screen()
+        self.assertTrue(self.screen.selected_disk is obj)
+    
+    def test_on_change_screen_do_copy(self):
+        '''Ensure disk is copied when do_copy flag is set'''
+        self.screen.install_profile = MockInstallProfile()
+        self.screen.install_profile.disk = True
+        
+        obj = object()
+        self.screen.disk_win.active_object = obj
+        self.screen.selected_disk = None
+        
+        disk = MockDiskInfo()
+        self.screen.disk_detail = MockAll()
+        self.screen.disk_detail.disk_info = disk
+        self.screen.do_copy = True
+        
+        self.screen.on_change_screen()
+        
+        self.assertTrue(self.screen.selected_disk is obj)
+        self.assertTrue(self.screen.install_profile.original_disk is disk)
+    
+    def test_on_change_screen_disk_is_none(self):
+        '''Check DiskScreen.on_change_screen when disk is None'''
+        self.screen.install_profile = MockInstallProfile()
+        
+        disk = MockDiskInfo()
+        self.screen.disk_detail = MockAll()
+        self.screen.disk_detail.disk_info = disk
+        
+        self.screen.on_change_screen()
+        
+        self.assertTrue(self.screen.install_profile.original_disk is disk)
+    
+    def test_size_line(self):
+        '''Ensure that DiskScreen._size_line is created and is a string after
+        calling get_size_line. Also verify that subsequent calls do not modify
+        the _size_line
+        
+        '''
+        self.assertTrue(self.screen._size_line is None)
+        self.screen.get_size_line()
+        self.assertTrue(isinstance(self.screen._size_line, basestring))
+        
+        obj = object()
+        self.screen._size_line = obj
+        self.screen.get_size_line()
+        self.assertTrue(obj is self.screen._size_line)
+    
+    def test_determine_size_data(self):
+        '''Ensure that recommended_size and minimum_size are accessible after
+        a call to determine_size_data(), and that they are numbers'''
+        
+        self.assertTrue(self.screen._recommended_size is None)
+        self.assertTrue(self.screen._minimum_size is None)
+        self.screen.determine_size_data()
+        self.assertTrue(isinstance(self.screen.minimum_size, numbers.Real))
+        self.assertTrue(isinstance(self.screen.recommended_size, numbers.Real))
+
+
+if __name__ == '__main__':
+    unittest.main()