Backend RPC handlers for Moblab OOBE related calls.

BUG=chromium:596536

TEST=Unit tested.

Change-Id: Iaa113708138f6a3fe88f44057d66a8faf240b04a
Reviewed-on: https://chromium-review.googlesource.com/334191
Commit-Ready: Michael Tang <ntang@chromium.org>
Tested-by: Michael Tang <ntang@chromium.org>
Reviewed-by: Dan Shi <dshi@chromium.org>
Reviewed-by: Michael Tang <ntang@chromium.org>
diff --git a/frontend/afe/site_rpc_interface.py b/frontend/afe/site_rpc_interface.py
index e709ba7..797ca06 100644
--- a/frontend/afe/site_rpc_interface.py
+++ b/frontend/afe/site_rpc_interface.py
@@ -6,12 +6,21 @@
 
 __author__ = 'cmasone@chromium.org (Chris Masone)'
 
+# The boto module is only available/used in Moblab for validation of cloud
+# storage access. The module is not available in the test lab environment,
+# and the import error is handled.
+try:
+    import boto
+except ImportError:
+    boto = None
 import common
 import ConfigParser
 import datetime
 import logging
 import os
+import re
 import shutil
+import socket
 
 from autotest_lib.frontend.afe import models
 from autotest_lib.client.common_lib import control_data
@@ -38,6 +47,20 @@
 _CONFIG = global_config.global_config
 MOBLAB_BOTO_LOCATION = '/home/moblab/.boto'
 
+# Google Cloud Storage bucket url regex pattern. The pattern is used to extract
+# the bucket name from the bucket URL. For example, "gs://image_bucket/google"
+# should result in a bucket name "image_bucket".
+GOOGLE_STORAGE_BUCKET_URL_PATTERN = re.compile(
+        r'gs://(?P<bucket>[a-zA-Z][a-zA-Z0-9-_]*)/?.*')
+
+# Constants used in JSON RPC field names.
+_USE_EXISTING_BOTO_FILE = 'use_existing_boto_file'
+_GS_ACCESS_KEY_ID = 'gs_access_key_id'
+_GS_SECRETE_ACCESS_KEY = 'gs_secret_access_key'
+_IMAGE_STORAGE_SERVER = 'image_storage_server'
+_RESULT_STORAGE_SERVER = 'results_storage_server'
+
+
 # Relevant CrosDynamicSuiteExceptions are defined in client/common_lib/error.py.
 
 
@@ -301,15 +324,88 @@
     return rpc_utils.prepare_for_serialization(config_values)
 
 
-@moblab_only
-def update_config_handler(config_values):
-    """
-    Update config values and override shadow config.
+def _write_config_file(config_file, config_values, overwrite=False):
+    """Writes out a configuration file.
 
-    @param config_values: See get_moblab_settings().
+    @param config_file: The name of the configuration file.
+    @param config_values: The ConfigParser object.
+    @param ovewrite: Flag on if overwriting is allowed.
+    """
+    if not config_file:
+        raise error.RPCException('Empty config file name.')
+    if not overwrite and os.path.exists(config_file):
+        raise error.RPCException('Config file already exists.')
+
+    if config_values:
+        with open(config_file, 'w') as config_file:
+            config_values.write(config_file)
+
+
+def _read_original_config():
+    """Reads the orginal configuratino without shadow.
+
+    @return: A configuration object, see global_config_class.
     """
     original_config = global_config.global_config_class()
     original_config.set_config_files(shadow_file='')
+    return original_config
+
+
+def _read_raw_config(config_file):
+    """Reads the raw configuration from a configuration file.
+
+    @param: config_file: The path of the configuration file.
+
+    @return: A ConfigParser object.
+    """
+    shadow_config = ConfigParser.RawConfigParser()
+    shadow_config.read(config_file)
+    return shadow_config
+
+
+def _get_shadow_config_from_partial_update(config_values):
+    """Finds out the new shadow configuration based on a partial update.
+
+    Since the input is only a partial config, we should not lose the config
+    data inside the existing shadow config file. We also need to distinguish
+    if the input config info overrides with a new value or reverts back to
+    an original value.
+
+    @param config_values: See get_moblab_settings().
+
+    @return: The new shadow configuration as ConfigParser object.
+    """
+    original_config = _read_original_config()
+    existing_shadow = _read_raw_config(_CONFIG.shadow_file)
+    for section, config_value_list in config_values.iteritems():
+        for key, value in config_value_list:
+            if original_config.get_config_value(section, key,
+                                                default='',
+                                                allow_blank=True) != value:
+                if not existing_shadow.has_section(section):
+                    existing_shadow.add_section(section)
+                existing_shadow.set(section, key, value)
+            elif existing_shadow.has_option(section, key):
+                existing_shadow.remove_option(section, key)
+    return existing_shadow
+
+
+def _update_partial_config(config_values):
+    """Updates the shadow configuration file with a partial config udpate.
+
+    @param config_values: See get_moblab_settings().
+    """
+    existing_config = _get_shadow_config_from_partial_update(config_values)
+    _write_config_file(_CONFIG.shadow_file, existing_config, True)
+
+
+@moblab_only
+def update_config_handler(config_values):
+    """Update config values and override shadow config.
+
+    @param config_values: See get_moblab_settings().
+    """
+    original_config = _read_original_config()
     new_shadow = ConfigParser.RawConfigParser()
     for section, config_value_list in config_values.iteritems():
         for key, value in config_value_list:
@@ -319,11 +415,11 @@
                 if not new_shadow.has_section(section):
                     new_shadow.add_section(section)
                 new_shadow.set(section, key, value)
+
     if not _CONFIG.shadow_file or not os.path.exists(_CONFIG.shadow_file):
         raise error.RPCException('Shadow config file does not exist.')
+    _write_config_file(_CONFIG.shadow_file, new_shadow, True)
 
-    with open(_CONFIG.shadow_file, 'w') as config_file:
-        new_shadow.write(config_file)
     # TODO (sbasi) crbug.com/403916 - Remove the reboot command and
     # instead restart the services that rely on the config values.
     os.system('sudo reboot')
@@ -363,6 +459,246 @@
     os.system('sudo restart moblab-devserver-init')
 
 
+###########Moblab Config Wizard RPCs #######################
+def _get_public_ip_address(socket_handle):
+    """Gets the public IP address.
+
+    Connects to Google DNS server using a socket and gets the preferred IP
+    address from the connection.
+
+    @param: socket_handle: a unix socket.
+
+    @return: public ip address as string.
+    """
+    try:
+        socket_handle.settimeout(1)
+        socket_handle.connect(('8.8.8.8', 53))
+        socket_name = socket_handle.getsockname()
+        if socket_name is not None:
+            logging.info('Got socket name from UDP socket.')
+            return socket_name[0]
+        logging.warn('Created UDP socket but with no socket_name.')
+    except socket.error:
+        logging.warn('Could not get socket name from UDP socket.')
+    return None
+
+
+def _get_network_info():
+    """Gets the network information.
+
+    TCP socket is used to test the connectivity. If there is no connectivity, try to
+    get the public IP with UDP socket.
+
+    @return: a tuple as (public_ip_address, connected_to_internet).
+    """
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    ip = _get_public_ip_address(s)
+    if ip is not None:
+        logging.info('Established TCP connection with well known server.')
+        return (ip, True)
+    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+    return (_get_public_ip_address(s), False)
+
+
+@moblab_only
+def get_network_info():
+    """Returns the server ip addresses, and if the server connectivity.
+
+    The server ip addresses as an array of strings, and the connectivity as a
+    flag.
+    """
+    network_info = {}
+    info = _get_network_info()
+    if info[0] is not None:
+        network_info['server_ips'] = [info[0]]
+    network_info['is_connected'] = info[1]
+
+    return rpc_utils.prepare_for_serialization(network_info)
+
+
+# Gets the boto configuration.
+def _get_boto_config():
+    """Reads the boto configuration from the boto file.
+
+    @return: Boto configuration as ConfigParser object.
+    """
+    boto_config = ConfigParser.ConfigParser()
+    boto_config.read(MOBLAB_BOTO_LOCATION)
+    return boto_config
+
+
+@moblab_only
+def get_cloud_storage_info():
+    """RPC handler to get the cloud storage access information.
+    """
+    cloud_storage_info = {}
+    value =_CONFIG.get_config_value('CROS', _IMAGE_STORAGE_SERVER)
+    if value is not None:
+        cloud_storage_info[_IMAGE_STORAGE_SERVER] = value
+    value =_CONFIG.get_config_value('CROS', _RESULT_STORAGE_SERVER)
+    if value is not None:
+        cloud_storage_info[_RESULT_STORAGE_SERVER] = value
+
+    boto_config = _get_boto_config()
+    sections = boto_config.sections()
+
+    if sections:
+        cloud_storage_info[_USE_EXISTING_BOTO_FILE] = True
+    else:
+        cloud_storage_info[_USE_EXISTING_BOTO_FILE] = False
+    if 'Credentials' in sections:
+        options = boto_config.options('Credentials')
+        if _GS_ACCESS_KEY_ID in options:
+            value = boto_config.get('Credentials', _GS_ACCESS_KEY_ID)
+            cloud_storage_info[_GS_ACCESS_KEY_ID] = value
+        if _GS_SECRETE_ACCESS_KEY in options:
+            value = boto_config.get('Credentials', _GS_SECRETE_ACCESS_KEY)
+            cloud_storage_info[_GS_SECRETE_ACCESS_KEY] = value
+
+    return rpc_utils.prepare_for_serialization(cloud_storage_info)
+
+
+def _get_bucket_name_from_url(bucket_url):
+    """Gets the bucket name from a bucket url.
+
+    @param: bucket_url: the bucket url string.
+    """
+    if bucket_url:
+        match = GOOGLE_STORAGE_BUCKET_URL_PATTERN.match(bucket_url)
+        if match:
+            return match.group('bucket')
+    return None
+
+
+def _is_valid_boto_key(key_id, key_secret):
+    """Checks if the boto key is valid.
+
+    @param: key_id: The boto key id string.
+    @param: key_secret: The boto key string.
+
+    @return: A tuple as (valid_boolean, details_string).
+    """
+    if not key_id or not key_secret:
+        return (False, "Empty key id or secret.")
+    conn = boto.connect_gs(key_id, key_secret)
+    try:
+        buckets = conn.get_all_buckets()
+        return (True, None)
+    except boto.exception.GSResponseError:
+        details = "The boto access key is not valid"
+        return (False, details)
+    finally:
+        conn.close()
+
+
+def _is_valid_bucket(key_id, key_secret, bucket_name):
+    """Checks if a bucket is valid and accessible.
+
+    @param: key_id: The boto key id string.
+    @param: key_secret: The boto key string.
+    @param: bucket name string.
+
+    @return: A tuple as (valid_boolean, details_string).
+    """
+    if not key_id or not key_secret or not bucket_name:
+        return (False, "Server error: invalid argument")
+    conn = boto.connect_gs(key_id, key_secret)
+    bucket = conn.lookup(bucket_name)
+    conn.close()
+    if bucket:
+        return (True, None)
+    return (False, "Bucket %s does not exist." % bucket_name)
+
+
+def _is_valid_bucket_url(key_id, key_secret, bucket_url):
+    """Validates the bucket url is accessible.
+
+    @param: key_id: The boto key id string.
+    @param: key_secret: The boto key string.
+    @param: bucket url string.
+
+    @return: A tuple as (valid_boolean, details_string).
+    """
+    bucket_name = _get_bucket_name_from_url(bucket_url)
+    if bucket_name:
+        return _is_valid_bucket(key_id, key_secret, bucket_name)
+    return (False, "Bucket url %s is not valid" % bucket_url)
+
+
+def _validate_cloud_storage_info(cloud_storage_info):
+    """Checks if the cloud storage information is valid.
+
+    @param: cloud_storage_info: The JSON RPC object for cloud storage info.
+
+    @return: A tuple as (valid_boolean, details_string).
+    """
+    valid = True
+    details = None
+    if not cloud_storage_info[_USE_EXISTING_BOTO_FILE]:
+        key_id = cloud_storage_info[_GS_ACCESS_KEY_ID]
+        key_secret = cloud_storage_info[_GS_SECRETE_ACCESS_KEY]
+        valid, details = _is_valid_boto_key(key_id, key_secret)
+
+        if valid:
+            valid, details = _is_valid_bucket_url(
+                key_id, key_secret, cloud_storage_info[_IMAGE_STORAGE_SERVER])
+
+        if valid:
+            valid, details = _is_valid_bucket_url(
+                key_id, key_secret, cloud_storage_info[_RESULT_STORAGE_SERVER])
+    return (valid, details)
+
+
+def _create_operation_status_response(is_ok, details):
+    """Helper method to create a operation status reponse.
+
+    @param: is_ok: Boolean for if the operation is ok.
+    @param: details: A detailed string.
+
+    @return: A serialized JSON RPC object.
+    """
+    status_response = {'status_ok': is_ok}
+    if details:
+        status_response['status_details'] = details
+    return rpc_utils.prepare_for_serialization(status_response)
+
+
+@moblab_only
+def validate_cloud_storage_info(cloud_storage_info):
+    """RPC handler to check if the cloud storage info is valid.
+    """
+    valid, details = _validate_cloud_storage_info(cloud_storage_info)
+    return _create_operation_status_response(valid, details)
+
+
+@moblab_only
+def submit_wizard_config_info(cloud_storage_info):
+    """RPC handler to submit the cloud storage info.
+    """
+    valid, details = _validate_cloud_storage_info(cloud_storage_info)
+    if not valid:
+        return _create_operation_status_response(valid, details)
+    config_update = {}
+    config_update['CROS'] = [
+        (_IMAGE_STORAGE_SERVER, cloud_storage_info[_IMAGE_STORAGE_SERVER]),
+        (_RESULT_STORAGE_SERVER, cloud_storage_info[_RESULT_STORAGE_SERVER])
+    ]
+    _update_partial_config(config_update)
+
+    if not cloud_storage_info[_USE_EXISTING_BOTO_FILE]:
+        boto_config = ConfigParser.RawConfigParser()
+        boto_config.add_section('Credentials')
+        boto_config.set('Credentials', _GS_ACCESS_KEY_ID,
+                        cloud_storage_info[_GS_ACCESS_KEY_ID])
+        boto_config.set('Credentials', _GS_SECRETE_ACCESS_KEY,
+                        cloud_storage_info[_GS_SECRETE_ACCESS_KEY])
+        _write_config_file(MOBLAB_BOTO_LOCATION, boto_config, True)
+
+    _CONFIG.parse_config_file()
+
+    return _create_operation_status_response(True, None)
+
+
 def get_job_history(**filter_data):
     """Get history of the job, including the special tasks executed for the job
 
diff --git a/frontend/afe/site_rpc_interface_unittest.py b/frontend/afe/site_rpc_interface_unittest.py
index db4a300..20330f0 100755
--- a/frontend/afe/site_rpc_interface_unittest.py
+++ b/frontend/afe/site_rpc_interface_unittest.py
@@ -8,7 +8,16 @@
 
 
 import __builtin__
+# The boto module is only available/used in Moblab for validation of cloud
+# storage access. The module is not available in the test lab environment, 
+# and the import error is handled.
+try:
+    import boto
+except ImportError:
+    boto = None
+import ConfigParser
 import datetime
+import logging
 import mox
 import StringIO
 import unittest
@@ -358,6 +367,275 @@
         site_rpc_interface.get_config_values()
 
 
+    def testGetNetworkInfo(self):
+        """Ensure the network info is properly converted to a dict."""
+        self.setIsMoblab(True)
+
+        self.mox.StubOutWithMock(site_rpc_interface, '_get_network_info')
+        site_rpc_interface._get_network_info().AndReturn(('10.0.0.1', True))
+        self.mox.StubOutWithMock(rpc_utils, 'prepare_for_serialization')
+
+        rpc_utils.prepare_for_serialization(
+               {'is_connected': True, 'server_ips': ['10.0.0.1']})
+        self.mox.ReplayAll()
+        site_rpc_interface.get_network_info()
+        self.mox.VerifyAll()
+
+
+    def testGetNetworkInfoWithNoIp(self):
+        """Queries network info with no public IP address."""
+        self.setIsMoblab(True)
+
+        self.mox.StubOutWithMock(site_rpc_interface, '_get_network_info')
+        site_rpc_interface._get_network_info().AndReturn((None, False))
+        self.mox.StubOutWithMock(rpc_utils, 'prepare_for_serialization')
+
+        rpc_utils.prepare_for_serialization(
+               {'is_connected': False})
+        self.mox.ReplayAll()
+        site_rpc_interface.get_network_info()
+        self.mox.VerifyAll()
+
+
+    def testGetNetworkInfoWithNoConnectivity(self):
+        """Queries network info with public IP address but no connectivity."""
+        self.setIsMoblab(True)
+
+        self.mox.StubOutWithMock(site_rpc_interface, '_get_network_info')
+        site_rpc_interface._get_network_info().AndReturn(('10.0.0.1', False))
+        self.mox.StubOutWithMock(rpc_utils, 'prepare_for_serialization')
+
+        rpc_utils.prepare_for_serialization(
+               {'is_connected': False, 'server_ips': ['10.0.0.1']})
+        self.mox.ReplayAll()
+        site_rpc_interface.get_network_info()
+        self.mox.VerifyAll()
+
+
+    def testGetCloudStorageInfo(self):
+        """Ensure the cloud storage info is properly converted to a dict."""
+        self.setIsMoblab(True)
+        config_mock = self.mox.CreateMockAnything()
+        site_rpc_interface._CONFIG = config_mock
+        config_mock.get_config_value(
+            'CROS', 'image_storage_server').AndReturn('gs://bucket1')
+        config_mock.get_config_value(
+            'CROS', 'results_storage_server').AndReturn('gs://bucket2')
+        self.mox.StubOutWithMock(site_rpc_interface, '_get_boto_config')
+        site_rpc_interface._get_boto_config().AndReturn(config_mock)
+        config_mock.sections().AndReturn(['Credentials', 'b'])
+        config_mock.options('Credentials').AndReturn(
+            ['gs_access_key_id', 'gs_secret_access_key'])
+        config_mock.get(
+            'Credentials', 'gs_access_key_id').AndReturn('key')
+        config_mock.get(
+            'Credentials', 'gs_secret_access_key').AndReturn('secret')
+        rpc_utils.prepare_for_serialization(
+                {
+                    'gs_access_key_id': 'key',
+                    'gs_secret_access_key' : 'secret',
+                    'use_existing_boto_file': True,
+                    'image_storage_server' : 'gs://bucket1',
+                    'results_storage_server' : 'gs://bucket2'
+                })
+        self.mox.ReplayAll()
+        site_rpc_interface.get_cloud_storage_info()
+        self.mox.VerifyAll()
+
+
+    def testValidateCloudStorageInfo(self):
+        """ Ensure the cloud storage info validation flow."""
+        self.setIsMoblab(True)
+        cloud_storage_info = {
+            'use_existing_boto_file': False,
+            'gs_access_key_id': 'key',
+            'gs_secret_access_key': 'secret',
+            'image_storage_server': 'gs://bucket1',
+            'results_storage_server': 'gs://bucket2'}
+        self.mox.StubOutWithMock(site_rpc_interface, '_is_valid_boto_key')
+        self.mox.StubOutWithMock(site_rpc_interface, '_is_valid_bucket')
+        site_rpc_interface._is_valid_boto_key(
+                'key', 'secret').AndReturn((True, None))
+        site_rpc_interface._is_valid_bucket(
+                'key', 'secret', 'bucket1').AndReturn((True, None))
+        site_rpc_interface._is_valid_bucket(
+                'key', 'secret', 'bucket2').AndReturn((True, None))
+        rpc_utils.prepare_for_serialization(
+                {'status_ok': True })
+        self.mox.ReplayAll()
+        site_rpc_interface.validate_cloud_storage_info(cloud_storage_info)
+        self.mox.VerifyAll()
+
+
+    def testGetBucketNameFromUrl(self):
+        """Gets bucket name from bucket URL."""
+        self.assertEquals(
+            'bucket_name-123',
+            site_rpc_interface._get_bucket_name_from_url(
+                    'gs://bucket_name-123'))
+        self.assertEquals(
+            'bucket_name-123',
+            site_rpc_interface._get_bucket_name_from_url(
+                    'gs://bucket_name-123/'))
+        self.assertEquals(
+            'bucket_name-123',
+            site_rpc_interface._get_bucket_name_from_url(
+                    'gs://bucket_name-123/a/b/c'))
+        self.assertIsNone(site_rpc_interface._get_bucket_name_from_url(
+            'bucket_name-123/a/b/c'))
+
+
+    def testIsValidBotoKeyValid(self):
+        """Tests the boto key validation flow."""
+        if boto is None:
+            logging.info('skip test since boto module not installed')
+            return
+        conn = self.mox.CreateMockAnything()
+        self.mox.StubOutWithMock(boto, 'connect_gs')
+        boto.connect_gs('key', 'secret').AndReturn(conn)
+        conn.get_all_buckets().AndReturn(['a', 'b'])
+        conn.close()
+        self.mox.ReplayAll()
+        valid, details = site_rpc_interface._is_valid_boto_key('key', 'secret')
+        self.assertTrue(valid)
+        self.mox.VerifyAll()
+
+
+    def testIsValidBotoKeyInvalid(self):
+        """Tests the boto key validation with invalid key."""
+        if boto is None:
+            logging.info('skip test since boto module not installed')
+            return
+        conn = self.mox.CreateMockAnything()
+        self.mox.StubOutWithMock(boto, 'connect_gs')
+        boto.connect_gs('key', 'secret').AndReturn(conn)
+        conn.get_all_buckets().AndRaise(
+                boto.exception.GSResponseError('bad', 'reason'))
+        conn.close()
+        self.mox.ReplayAll()
+        valid, details = site_rpc_interface._is_valid_boto_key('key', 'secret')
+        self.assertFalse(valid)
+        self.assertEquals('The boto access key is not valid', details)
+        self.mox.VerifyAll()
+
+
+    def testIsValidBucketValid(self):
+        """Tests the bucket vaildation flow."""
+        if boto is None:
+            logging.info('skip test since boto module not installed')
+            return
+        conn = self.mox.CreateMockAnything()
+        self.mox.StubOutWithMock(boto, 'connect_gs')
+        boto.connect_gs('key', 'secret').AndReturn(conn)
+        conn.lookup('bucket').AndReturn('bucket')
+        conn.close()
+        self.mox.ReplayAll()
+        valid, details = site_rpc_interface._is_valid_bucket(
+                'key', 'secret', 'bucket')
+        self.assertTrue(valid)
+        self.mox.VerifyAll()
+
+
+    def testIsValidBucketInvalid(self):
+        """Tests the bucket validation flow with invalid key."""
+        if boto is None:
+            logging.info('skip test since boto module not installed')
+            return
+        conn = self.mox.CreateMockAnything()
+        self.mox.StubOutWithMock(boto, 'connect_gs')
+        boto.connect_gs('key', 'secret').AndReturn(conn)
+        conn.lookup('bucket').AndReturn(None)
+        conn.close()
+        self.mox.ReplayAll()
+        valid, details = site_rpc_interface._is_valid_bucket(
+                'key', 'secret', 'bucket')
+        self.assertFalse(valid)
+        self.assertEquals("Bucket bucket does not exist.", details)
+        self.mox.VerifyAll()
+
+
+    def testGetShadowConfigFromPartialUpdate(self):
+        """Tests getting shadow configuration based on partial upate."""
+        partial_config = {
+                'section1': [
+                    ('opt1', 'value1'),
+                    ('opt2', 'value2'),
+                    ('opt3', 'value3'),
+                    ('opt4', 'value4'),
+                    ]
+                }
+        shadow_config_str = "[section1]\nopt2 = value2_1\nopt4 = value4_1"
+        shadow_config = ConfigParser.ConfigParser()
+        shadow_config.readfp(StringIO.StringIO(shadow_config_str))
+        original_config = self.mox.CreateMockAnything()
+        self.mox.StubOutWithMock(site_rpc_interface, '_read_original_config')
+        self.mox.StubOutWithMock(site_rpc_interface, '_read_raw_config')
+        site_rpc_interface._read_original_config().AndReturn(original_config)
+        site_rpc_interface._read_raw_config(
+                site_rpc_interface._CONFIG.shadow_file).AndReturn(shadow_config)
+        original_config.get_config_value(
+                'section1', 'opt1',
+                allow_blank=True, default='').AndReturn('value1')
+        original_config.get_config_value(
+                'section1', 'opt2',
+                allow_blank=True, default='').AndReturn('value2')
+        original_config.get_config_value(
+                'section1', 'opt3',
+                allow_blank=True, default='').AndReturn('blah')
+        original_config.get_config_value(
+                'section1', 'opt4',
+                allow_blank=True, default='').AndReturn('blah')
+        self.mox.ReplayAll()
+        shadow_config = site_rpc_interface._get_shadow_config_from_partial_update(
+                partial_config)
+        # opt1 same as the original.
+        self.assertFalse(shadow_config.has_option('section1', 'opt1'))
+        # opt2 reverts back to original
+        self.assertFalse(shadow_config.has_option('section1', 'opt2'))
+        # opt3 is updated from original.
+        self.assertEquals('value3', shadow_config.get('section1', 'opt3'))
+        # opt3 in shadow but updated again.
+        self.assertEquals('value4', shadow_config.get('section1', 'opt4'))
+        self.mox.VerifyAll()
+
+
+    def testGetShadowConfigFromPartialUpdateWithNewSection(self):
+        """
+        Test getting shadown configuration based on partial update with new section.
+        """
+        partial_config = {
+                'section2': [
+                    ('opt5', 'value5'),
+                    ('opt6', 'value6'),
+                    ],
+                }
+        shadow_config_str = "[section1]\nopt2 = value2_1\n"
+        shadow_config = ConfigParser.ConfigParser()
+        shadow_config.readfp(StringIO.StringIO(shadow_config_str))
+        original_config = self.mox.CreateMockAnything()
+        self.mox.StubOutWithMock(site_rpc_interface, '_read_original_config')
+        self.mox.StubOutWithMock(site_rpc_interface, '_read_raw_config')
+        site_rpc_interface._read_original_config().AndReturn(original_config)
+        site_rpc_interface._read_raw_config(
+            site_rpc_interface._CONFIG.shadow_file).AndReturn(shadow_config)
+        original_config.get_config_value(
+                'section2', 'opt5',
+                allow_blank=True, default='').AndReturn('value5')
+        original_config.get_config_value(
+                'section2', 'opt6',
+                allow_blank=True, default='').AndReturn('blah')
+        self.mox.ReplayAll()
+        shadow_config = site_rpc_interface._get_shadow_config_from_partial_update(
+                partial_config)
+        # opt2 is still in shadow
+        self.assertEquals('value2_1', shadow_config.get('section1', 'opt2'))
+        # opt5 is not changed.
+        self.assertFalse(shadow_config.has_option('section2', 'opt5'))
+        # opt6 is updated.
+        self.assertEquals('value6', shadow_config.get('section2', 'opt6'))
+        self.mox.VerifyAll()
+
+
     def _mockReadFile(self, path, lines=[]):
         """Mock out reading a file line by line.