AIDEGen: add test_mapping file type in IntelliJ.

Bug: 151700770
Test: 1. m aidegen; aidegen-dev Settings -s
      2. new a TEST_MAPPING file in IntelliJ, IntelliJ should recognize
         it as JSON file.

Change-Id: Id3994e1855aa2c44d13d95f52882a03c1fa165ae
diff --git a/aidegen/lib/ide_util.py b/aidegen/lib/ide_util.py
index 12c38ac..4569c9b 100644
--- a/aidegen/lib/ide_util.py
+++ b/aidegen/lib/ide_util.py
@@ -37,6 +37,8 @@
 import re
 import subprocess
 
+from xml.etree import ElementTree
+
 from aidegen import constant
 from aidegen import templates
 from aidegen.lib import aidegen_metrics
@@ -48,6 +50,7 @@
 from aidegen.lib import project_config
 from aidegen.lib import project_file_gen
 from aidegen.sdk import jdk_table
+from aidegen.lib import xml_util
 
 # Add 'nohup' to prevent IDE from being terminated when console is terminated.
 _IDEA_FOLDER = '.idea'
@@ -76,11 +79,20 @@
 LINUX_JDK_PATH = os.path.join(common_util.get_android_root_dir(),
                               'prebuilts/jdk/jdk8/linux-x86')
 LINUX_JDK_TABLE_PATH = 'config/options/jdk.table.xml'
+LINUX_FILE_TYPE_PATH = 'config/options/filetypes.xml'
 LINUX_ANDROID_SDK_PATH = os.path.join(os.getenv('HOME'), 'Android/Sdk')
 MAC_JDK_PATH = os.path.join(common_util.get_android_root_dir(),
                             'prebuilts/jdk/jdk8/darwin-x86')
 MAC_JDK_TABLE_PATH = 'options/jdk.table.xml'
+MAC_FILE_TYPE_XML_PATH = 'options/filetypes.xml'
 MAC_ANDROID_SDK_PATH = os.path.join(os.getenv('HOME'), 'Library/Android/sdk')
+PATTERN_KEY = 'pattern'
+TYPE_KEY = 'type'
+JSON_TYPE = 'JSON'
+TEST_MAPPING_NAME = 'TEST_MAPPING'
+_TEST_MAPPING_TYPE = '<mapping pattern="TEST_MAPPING" type="JSON" />'
+_XPATH_EXTENSION_MAP = 'component/extensionMap'
+_XPATH_MAPPING = _XPATH_EXTENSION_MAP + '/mapping'
 
 
 # pylint: disable=too-many-lines
@@ -146,6 +158,7 @@
     Class Attributes:
         _JDK_PATH: The path of JDK in android project.
         _IDE_JDK_TABLE_PATH: The path of JDK table which record JDK info in IDE.
+        _IDE_FILE_TYPE_PATH: The path of filetypes.xml.
         _JDK_CONTENT: A string, the content of the JDK configuration.
         _DEFAULT_ANDROID_SDK_PATH: A string, the path of Android SDK.
         _CONFIG_DIR: A string of the config folder name.
@@ -171,6 +184,7 @@
 
     _JDK_PATH = ''
     _IDE_JDK_TABLE_PATH = ''
+    _IDE_FILE_TYPE_PATH = ''
     _JDK_CONTENT = ''
     _DEFAULT_ANDROID_SDK_PATH = ''
     _CONFIG_DIR = ''
@@ -225,6 +239,45 @@
             intellij_config_dir = os.path.join(_config_path, self._CONFIG_DIR)
             config.IdeaProperties(intellij_config_dir).set_max_file_size()
 
+            self._add_test_mapping_file_type(_config_path)
+
+    def _add_test_mapping_file_type(self, _config_path):
+        """Adds TEST_MAPPING file type.
+
+        IntelliJ can't recognize TEST_MAPPING files as the json file. It needs
+        adding file type mapping in filetypes.xml to recognize TEST_MAPPING
+        files.
+
+        Args:
+            _config_path: the path of IDE config.
+        """
+        file_type_path = os.path.join(_config_path, self._IDE_FILE_TYPE_PATH)
+        if not os.path.isfile(file_type_path):
+            logging.warning('Filetypes.xml is not found.')
+            return
+
+        file_type_xml = xml_util.parse_xml(file_type_path)
+        if not file_type_xml:
+            logging.warning('Can\'t parse filetypes.xml.')
+            return
+
+        root = file_type_xml.getroot()
+        add_pattern = True
+        for mapping in root.findall(_XPATH_MAPPING):
+            attrib = mapping.attrib
+            if PATTERN_KEY in attrib and TYPE_KEY in attrib:
+                if attrib[PATTERN_KEY] == TEST_MAPPING_NAME:
+                    if attrib[TYPE_KEY] != JSON_TYPE:
+                        attrib[TYPE_KEY] = JSON_TYPE
+                        file_type_xml.write(file_type_path)
+                    add_pattern = False
+                    break
+        if add_pattern:
+            root.find(_XPATH_EXTENSION_MAP).append(
+                ElementTree.fromstring(_TEST_MAPPING_TYPE))
+            pretty_xml = common_util.to_pretty_xml(root)
+            common_util.file_generate(file_type_path, pretty_xml)
+
     def _get_config_root_paths(self):
         """Get the config root paths from derived class.
 
@@ -565,6 +618,7 @@
     # TODO(b/127899277): Preserve a config for jdk version option case.
     _CONFIG_DIR = CONFIG_DIR
     _IDE_JDK_TABLE_PATH = LINUX_JDK_TABLE_PATH
+    _IDE_FILE_TYPE_PATH = LINUX_FILE_TYPE_PATH
     _JDK_CONTENT = templates.LINUX_JDK_XML
     _DEFAULT_ANDROID_SDK_PATH = LINUX_ANDROID_SDK_PATH
     _SYMBOLIC_VERSIONS = ['/opt/intellij-ce-stable/bin/idea.sh',
@@ -636,6 +690,7 @@
 
     _JDK_PATH = MAC_JDK_PATH
     _IDE_JDK_TABLE_PATH = MAC_JDK_TABLE_PATH
+    _IDE_FILE_TYPE_PATH = MAC_FILE_TYPE_XML_PATH
     _JDK_CONTENT = templates.MAC_JDK_XML
     _DEFAULT_ANDROID_SDK_PATH = MAC_ANDROID_SDK_PATH
 
diff --git a/aidegen/lib/ide_util_unittest.py b/aidegen/lib/ide_util_unittest.py
index b055caf..9a59360 100644
--- a/aidegen/lib/ide_util_unittest.py
+++ b/aidegen/lib/ide_util_unittest.py
@@ -24,10 +24,12 @@
 import unittest
 
 from unittest import mock
+from xml.etree import ElementTree
 
 from aidegen import aidegen_main
 from aidegen import unittest_constants
 from aidegen.lib import android_dev_os
+from aidegen.lib import common_util
 from aidegen.lib import config
 from aidegen.lib import errors
 from aidegen.lib import ide_common_util
@@ -40,6 +42,7 @@
 # pylint: disable=too-many-public-methods
 # pylint: disable=protected-access
 # pylint: disable-msg=too-many-arguments
+# pylint: disable-msg=unused-argument
 class IdeUtilUnittests(unittest.TestCase):
     """Unit tests for ide_util.py."""
 
@@ -49,6 +52,22 @@
     _TEST_PRJ_PATH4 = ''
     _MODULE_XML_SAMPLE = ''
     _TEST_DIR = None
+    _TEST_XML_CONTENT = """<application>
+  <component name="FileTypeManager" version="17">
+    <extensionMap>
+      <mapping ext="pi" type="Python"/>
+    </extensionMap>
+  </component>
+</application>"""
+    _TEST_XML_CONTENT_2 = """<application>
+  <component name="FileTypeManager" version="17">
+    <extensionMap>
+      <mapping ext="pi" type="Python"/>
+      <mapping pattern="test" type="a"/>
+      <mapping pattern="TEST_MAPPING" type="a"/>
+    </extensionMap>
+  </component>
+</application>"""
 
     def setUp(self):
         """Prepare the testdata related path."""
@@ -220,12 +239,14 @@
         ide_obj._get_config_root_paths()
         self.assertTrue(mock_filter.called)
 
+    @mock.patch.object(ide_util.IdeBase, '_add_test_mapping_file_type')
     @mock.patch.object(config.IdeaProperties, 'set_max_file_size')
     @mock.patch.object(project_file_gen, 'gen_enable_debugger_module')
     @mock.patch.object(jdk_table.JDKTableXML, 'config_jdk_table_xml')
     @mock.patch.object(ide_util.IdeBase, '_get_config_root_paths')
     def test_apply_optional_config(self, mock_path, mock_config_xml,
-                                   mock_gen_debugger, mock_set_size):
+                                   mock_gen_debugger, mock_set_size,
+                                   mock_test_mapping):
         """Test basic logic of apply_optional_config."""
         ide = ide_util.IdeBase()
         ide._installed_path = None
@@ -248,6 +269,34 @@
         self.assertTrue(mock_gen_debugger.called)
         self.assertTrue(mock_set_size.called)
 
+    @mock.patch('os.path.isfile')
+    @mock.patch.object(ElementTree.ElementTree, 'write')
+    @mock.patch.object(common_util, 'to_pretty_xml')
+    @mock.patch.object(common_util, 'file_generate')
+    @mock.patch.object(ElementTree, 'parse')
+    @mock.patch.object(ElementTree.ElementTree, 'getroot')
+    def test_add_test_mapping_file_type(self, mock_root, mock_parse,
+                                        mock_file_gen, mock_pretty_xml,
+                                        mock_write, mock_isfile):
+        """Test basic logic of _add_test_mapping_file_type."""
+        mock_isfile.return_value = False
+        self.assertFalse(mock_file_gen.called)
+
+        mock_isfile.return_value = True
+        mock_parse.return_value = None
+        self.assertFalse(mock_file_gen.called)
+
+        mock_parse.return_value = ElementTree.ElementTree()
+        mock_root.return_value = ElementTree.fromstring(
+            self._TEST_XML_CONTENT_2)
+        ide_obj = ide_util.IdeBase()
+        ide_obj._add_test_mapping_file_type('')
+        self.assertFalse(mock_file_gen.called)
+        mock_root.return_value = ElementTree.fromstring(self._TEST_XML_CONTENT)
+        ide_obj._add_test_mapping_file_type('')
+        self.assertTrue(mock_pretty_xml.called)
+        self.assertTrue(mock_file_gen.called)
+
     @mock.patch('os.path.realpath')
     @mock.patch('os.path.isfile')
     def test_merge_symbolic_version(self, mock_isfile, mock_realpath):
@@ -417,12 +466,14 @@
         version = ide_intellij._get_preferred_version()
         self.assertEqual(version, '/a/b')
 
+    @mock.patch.object(ide_util.IdeBase, '_add_test_mapping_file_type')
     @mock.patch.object(ide_common_util, 'ask_preference')
     @mock.patch.object(config.IdeaProperties, 'set_max_file_size')
     @mock.patch.object(project_file_gen, 'gen_enable_debugger_module')
     @mock.patch.object(ide_util.IdeStudio, '_get_config_root_paths')
     def test_android_studio_class(self, mock_get_config_paths,
-                                  mock_gen_debugger, mock_set_size, mock_ask):
+                                  mock_gen_debugger, mock_set_size, mock_ask,
+                                  mock_add_file_type):
         """Test IdeStudio."""
         mock_get_config_paths.return_value = ['path1', 'path2']
         mock_gen_debugger.return_value = True