8000 add manage dataset and tests · dalequark/python-docs-samples@0f9db84 · GitHub
[go: up one dir, main page]

Skip to content

Commit 0f9db84

Browse files
committed
add manage dataset and tests
1 parent 3ffc66d commit 0f9db84

File tree

2 files changed

+54
-27
lines changed

2 files changed

+54
-27
lines changed

datalabeling/manage_dataset.py

Lines changed: 26 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,9 @@ def create_dataset(project_id):
2626
formatted_project_name = client.project_path(project_id)
2727

2828
dataset = datalabeling.types.Dataset(
29-
display_name='YOUR_ANNOTATION_SPEC_SET_DISPLAY_NAME',
30-
description='YOUR_DESCRIPTION')
29+
display_name='YOUR_ANNOTATION_SPEC_SET_DISPLAY_NAME',
30+
description='YOUR_DESCRIPTION'
31+
)
3132

3233
response = client.create_dataset(formatted_project_name, dataset)
3334

@@ -93,31 +94,41 @@ def delete_dataset(dataset_resource_name):
9394

9495
if __name__ == '__main__':
9596
parser = argparse.ArgumentParser(
96-
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
97+
description=__doc__,
98+
formatter_class=argparse.RawDescriptionHelpFormatter
99+
)
97100

98101
subparsers = parser.add_subparsers(dest='command')
99102

100103
create_parser = subparsers.add_parser('create', help='Create a new dataset.')
101104
create_parser.add_argument(
102-
'--project-id', help='Project ID. Required.', required=True)
105+
'--project-id',
106+
help='Project ID. Required.',
107+
required=True
108+
)
103109

104110
list_parser = subparsers.add_parser('list', help='List all datasets.')
105111
list_parser.add_argument(
106-
'--project-id', help='Project ID. Required.', required=True)
112+
'--project-id',
113+
help='Project ID. Required.',
114+
required=True
115+
)
107116

108-
get_parser = subparsers.add_parser(
109-
'get', help='Get a dataset by the dataset resource name.')
117+
get_parser = subparsers.add_parser('get',
118+
help='Get a dataset by the dataset resource name.')
110119
get_parser.add_argument(
111-
'--dataset-resource-name',
112-
help='The dataset resource name. Used in the get or delete operation.',
113-
required=True)
120+
'--dataset-resource-name',
121+
help='The dataset resource name. Used in the get or delete operation.',
122+
required=True
123+
)
114124

115-
delete_parser = subparsers.add_parser(
116-
'delete', help='Delete a dataset by the dataset resource name.')
125+
delete_parser = subparsers.add_parser('delete',
126+
help='Delete a dataset by the dataset resource name.')
117127
delete_parser.add_argument(
118-
'--dataset-resource-name',
119-
help='The dataset resource name. Used in the get or delete operation.',
120-
required=True)
128+
'--dataset-resource-name',
129+
help='The dataset resource name. Used in the get or delete operation.',
130+
required=True
131+
)
121132

122133
args = parser.parse_args()
123134

datalabeling/manage_dataset_test.py

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -22,26 +22,42 @@
2222
PROJECT_ID = os.getenv("GCLOUD_PROJECT")
2323

2424

25-
@pytest.mark.slow
26-
def test_manage_dataset(capsys):
27-
# Creates a dataset.
28-
manage_dataset.create_dataset(PROJECT_ID)
25+
@pytest.fixture(scope='function')
26+
def dataset():
27+
# create a temporary dataset
28+
dataset = manage_dataset.create_dataset(PROJECT_ID)
29+
30+
yield dataset
31+
32+
# tear down
33+
manage_dataset.delete_dataset(dataset.name)
34+
35+
36+
def test_create_dataset(capsys):
37+
response = manage_dataset.create_dataset(PROJECT_ID)
2938
out, _ = capsys.readouterr()
3039
assert "The dataset resource name:" in out
31-
create_dataset_output = out.splitlines()
32-
created_dataset_resource_name = create_dataset_output[0].split()[4]
3340

34-
# Lists the datasets
41+
# clean up
42+
manage_dataset.delete_dataset(response.name)
43+
44+
45+
def test_list_dataset(capsys, dataset):
3546
manage_dataset.list_datasets(PROJECT_ID)
3647
out, _ = capsys.readouterr()
37-
assert "The dataset resource name:" in out
48+
assert dataset.name in out
49+
3850

39-
# Gets the dataset.
40-
manage_dataset.get_dataset(created_dataset_resource_name)
51+
def test_get_dataset(capsys, dataset):
52+
manage_dataset.get_dataset(dataset.name)
4153
out, _ = capsys.readouterr()
4254
assert "The dataset resource name:" in out
4355

44-
# Deletes the dataset.
45-
manage_dataset.delete_dataset(created_dataset_resource_name)
56+
57+
def test_delete_dataset(capsys):
58+
# Creates a dataset.
59+
response = manage_dataset.create_dataset(PROJECT_ID)
60+
61+
manage_dataset.delete_dataset(response.name)
4662
out, _ = capsys.readouterr()
4763
assert "Dataset deleted." in out

0 commit comments

Comments
 (0)
0