Update DLP metadata samples and fix DLP quickstart

Author: andrewsg

dlp/metadata.py

@@ -21,79 +21,43 @@
 
 
 # [START list_info_types]
-def list_info_types(category, language_code='en-US'):
+def list_info_types(language_code=None, result_filter=None):
     """List types of sensitive information within a category.
     Args:
-        category: The category of info types to list; e.g. 'PII'.
         language_code: The BCP-47 language code to use, e.g. 'en-US'.
+        filter: An optional filter to only return info types supported by 
+                certain parts of the API. Defaults to "supported_by=INSPECT".
     Returns:
         None; the response from the API is printed to the terminal.
     """
     # Import the client library
-    import google.cloud.dlp_v2beta1
+    import google.cloud.dlp
 
     # Instantiate a client.
-    dlp = google.cloud.dlp_v2beta1.DlpServiceClient()
+    dlp = google.cloud.dlp.DlpServiceClient()
 
     # Make the API call.
-    response = dlp.list_info_types(category, language_code)
+    response = dlp.list_info_types(language_code, result_filter)
 
     # Print the results to the console.
-    print('Info types in {category}:'.format(category=category))
+    print('Info types:')
     for info_type in response.info_types:
         print('{name}: {display_name}'.format(
             name=info_type.name, display_name=info_type.display_name))
 # [END list_info_types]
 
 
-# [START list_categories]
-def list_categories(language_code='en-US'):
-    """List root categories of sensitive information.
-    Args:
-        language_code: The BCP-47 language code to use, e.g. 'en-US'.
-    Returns:
-        None; the response from the API is printed to the terminal.
-    """
-    # Import the client library
-    import google.cloud.dlp_v2beta1
-
-    # Instantiate a client.
-    dlp = google.cloud.dlp_v2beta1.DlpServiceClient()
-
-    # Make the API call.
-    response = dlp.list_root_categories(language_code)
-
-    # Print the results to the console.
-    print('Categories:')
-    for category in response.categories:
-        print('{name}: {display_name}'.format(
-            name=category.name, display_name=category.display_name))
-# [END list_categories]
-
-
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description=__doc__)
-    subparsers = parser.add_subparsers(
-        dest='metadata', help='Select which type of metadata to view.')
-
-    parser_categories = subparsers.add_parser(
-        'categories', help='Fetch the list of info type categories.')
-    parser_categories.add_argument(
-        '--language_code',
-        help='The BCP-47 language code to use, e.g. \'en-US\'.')
-
-    parser_info_types = subparsers.add_parser(
-        'info_types',
-        help='Fetch the list of info types in a specified category.')
-    parser_info_types.add_argument(
-        'category', help='The category of info types to list; e.g. \'PII\'.')
-    parser_info_types.add_argument(
+    parser.add_argument(
         '--language_code',
         help='The BCP-47 language code to use, e.g. \'en-US\'.')
+    parser.add_argument(
+        '--filter',
+        help='An optional filter to only return info types supported by '
+             'certain parts of the API. Defaults to "supported_by=INSPECT".')
 
     args = parser.parse_args()
 
-    if args.metadata == 'categories':
-        list_categories(language_code=args.language_code)
-    elif args.metadata == 'info_types':
-        list_info_types(args.category, language_code=args.language_code)
+    list_info_types(
+        language_code=args.language_code, result_filter=args.filter)

dlp/metadata_test.py

@@ -15,15 +15,8 @@
 import metadata
 
 
-def test_fetch_categories(capsys):
-    metadata.list_categories()
-
-    out, _ = capsys.readouterr()
-    assert 'PII' in out
-
-
 def test_fetch_info_types(capsys):
-    metadata.list_info_types('PII')
+    metadata.list_info_types()
 
     out, _ = capsys.readouterr()
     assert 'EMAIL_ADDRESS' in out

dlp/quickstart_test.py

@@ -12,11 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import mock
+import os
+
+import google.cloud.dlp
+
 import quickstart
 
 
+GCLOUD_PROJECT = os.getenv('GCLOUD_PROJECT')
+
 def test_quickstart(capsys):
-    quickstart.quickstart()
+    # Mock out project_path to use the test runner's project ID.
+    with mock.patch.object(
+            google.cloud.dlp.DlpServiceClient,
+            'project_path',
+            return_value='projects/{}'.format(GCLOUD_PROJECT)):
+        quickstart.quickstart()
 
     out, _ = capsys.readouterr()
-    assert 'US_MALE_NAME' in out
+    assert 'FIRST_NAME' in out
+    assert 'LAST_NAME' in out