|
22 | 22 | from google.api import client_pb2
|
23 | 23 | from google.api import resource_pb2
|
24 | 24 | from google.api_core import exceptions
|
| 25 | +from google.cloud import extended_operations_pb2 as ex_ops_pb2 |
25 | 26 | from google.gapic.metadata import gapic_metadata_pb2
|
26 | 27 | from google.longrunning import operations_pb2
|
27 | 28 | from google.protobuf import descriptor_pb2
|
@@ -1595,3 +1596,151 @@ def test_http_options(fs):
|
1595 | 1596 | method='get', uri='/v3/{name=projects/*/locations/*/operations/*}', body=None),
|
1596 | 1597 | wrappers.HttpRule(method='get', uri='/v3/{name=/locations/*/operations/*}', body=None)]
|
1597 | 1598 | }
|
| 1599 | + |
| 1600 | + |
| 1601 | +def generate_basic_extended_operations_setup(): |
| 1602 | + T = descriptor_pb2.FieldDescriptorProto.Type |
| 1603 | + |
| 1604 | + operation = make_message_pb2( |
| 1605 | + name="Operation", |
| 1606 | + fields=( |
| 1607 | + make_field_pb2(name=name, type=T.Value("TYPE_STRING"), number=i) |
| 1608 | + for i, name in enumerate(("name", "status", "error_code", "error_message"), start=1) |
| 1609 | + ), |
| 1610 | + ) |
| 1611 | + |
| 1612 | + for f in operation.field: |
| 1613 | + options = descriptor_pb2.FieldOptions() |
| 1614 | + # Note: The field numbers were carefully chosen to be the corresponding enum values. |
| 1615 | + options.Extensions[ex_ops_pb2.operation_field] = f.number |
| 1616 | + f.options.MergeFrom(options) |
| 1617 | + |
| 1618 | + options = descriptor_pb2.MethodOptions() |
| 1619 | + options.Extensions[ex_ops_pb2.operation_polling_method] = True |
| 1620 | + |
| 1621 | + polling_method = descriptor_pb2.MethodDescriptorProto( |
| 1622 | + name="Get", |
| 1623 | + input_type="google.extended_operations.v1.stuff.GetOperation", |
| 1624 | + output_type="google.extended_operations.v1.stuff.Operation", |
| 1625 | + options=options, |
| 1626 | + ) |
| 1627 | + |
| 1628 | + delete_input_message = make_message_pb2(name="Input") |
| 1629 | + delete_output_message = make_message_pb2(name="Output") |
| 1630 | + ops_service = descriptor_pb2.ServiceDescriptorProto( |
| 1631 | + name="CustomOperations", |
| 1632 | + method=[ |
| 1633 | + polling_method, |
| 1634 | + descriptor_pb2.MethodDescriptorProto( |
| 1635 | + name="Delete", |
| 1636 | + input_type="google.extended_operations.v1.stuff.Input", |
| 1637 | + output_type="google.extended_operations.v1.stuff.Output", |
| 1638 | + ), |
| 1639 | + ], |
| 1640 | + ) |
| 1641 | + |
| 1642 | + request = make_message_pb2( |
| 1643 | + name="GetOperation", |
| 1644 | + fields=[ |
| 1645 | + make_field_pb2(name="name", type=T.Value("TYPE_STRING"), number=1) |
| 1646 | + ], |
| 1647 | + ) |
| 1648 | + |
| 1649 | + initial_opts = descriptor_pb2.MethodOptions() |
| 1650 | + initial_opts.Extensions[ex_ops_pb2.operation_service] = ops_service.name |
| 1651 | + initial_input_message = make_message_pb2(name="Initial") |
| 1652 | + initial_method = descriptor_pb2.MethodDescriptorProto( |
| 1653 | + name="CreateTask", |
| 1654 | + input_type="google.extended_operations.v1.stuff.GetOperation", |
| 1655 | + output_type="google.extended_operations.v1.stuff.Operation", |
| 1656 | + options=initial_opts, |
| 1657 | + ) |
| 1658 | + |
| 1659 | + regular_service = descriptor_pb2.ServiceDescriptorProto( |
| 1660 | + name="RegularService", |
| 1661 | + method=[ |
| 1662 | + initial_method, |
| 1663 | + ], |
| 1664 | + ) |
| 1665 | + |
| 1666 | + file_protos = [ |
| 1667 | + make_file_pb2( |
| 1668 | + name="extended_operations.proto", |
| 1669 | + package="google.extended_operations.v1.stuff", |
| 1670 | + messages=[ |
| 1671 | + operation, |
| 1672 | + request, |
| 1673 | + delete_output_message, |
| 1674 | + delete_input_message, |
| 1675 | + initial_input_message, |
| 1676 | + ], |
| 1677 | + services=[ |
| 1678 | + regular_service, |
| 1679 | + ops_service, |
| 1680 | + ], |
| 1681 | + ), |
| 1682 | + ] |
| 1683 | + |
| 1684 | + return file_protos |
| 1685 | + |
| 1686 | + |
| 1687 | +def test_extended_operations_lro_operation_service(): |
| 1688 | + file_protos = generate_basic_extended_operations_setup() |
| 1689 | + api_schema = api.API.build(file_protos) |
| 1690 | + initial_method = api_schema.services["google.extended_operations.v1.stuff.RegularService"].methods["CreateTask"] |
| 1691 | + |
| 1692 | + expected = api_schema.services['google.extended_operations.v1.stuff.CustomOperations'] |
| 1693 | + actual = api_schema.get_custom_operation_service(initial_method) |
| 1694 | + |
| 1695 | + assert expected is actual |
| 1696 | + |
| 1697 | + assert actual.custom_polling_method is actual.methods["Get"] |
| 1698 | + |
| 1699 | + |
| 1700 | +def test_extended_operations_lro_operation_service_no_annotation(): |
| 1701 | + file_protos = generate_basic_extended_operations_setup() |
| 1702 | + |
| 1703 | + api_schema = api.API.build(file_protos) |
| 1704 | + initial_method = api_schema.services["google.extended_operations.v1.stuff.RegularService"].methods["CreateTask"] |
| 1705 | + # It's easier to manipulate data structures after building the API. |
| 1706 | + del initial_method.options.Extensions[ex_ops_pb2.operation_service] |
| 1707 | + |
| 1708 | + with pytest.raises(KeyError): |
| 1709 | + api_schema.get_custom_operation_service(initial_method) |
| 1710 | + |
| 1711 | + |
| 1712 | +def test_extended_operations_lro_operation_service_no_such_service(): |
| 1713 | + file_protos = generate_basic_extended_operations_setup() |
| 1714 | + |
| 1715 | + api_schema = api.API.build(file_protos) |
| 1716 | + initial_method = api_schema.services["google.extended_operations.v1.stuff.RegularService"].methods["CreateTask"] |
| 1717 | + initial_method.options.Extensions[ex_ops_pb2.operation_service] = "UnrealService" |
| 1718 | + |
| 1719 | + with pytest.raises(KeyError): |
| 1720 | + api_schema.get_custom_operation_service(initial_method) |
| 1721 | + |
| 1722 | + |
| 1723 | +def test_extended_operations_lro_operation_service_not_an_lro(): |
| 1724 | + file_protos = generate_basic_extended_operations_setup() |
| 1725 | + |
| 1726 | + api_schema = api.API.build(file_protos) |
| 1727 | + initial_method = api_schema.services["google.extended_operations.v1.stuff.RegularService"].methods["CreateTask"] |
| 1728 | + # Hack to pretend that the initial_method is not an LRO |
| 1729 | + super(type(initial_method), initial_method).__setattr__( |
| 1730 | + "output", initial_method.input) |
| 1731 | + |
| 1732 | + with pytest.raises(ValueError): |
| 1733 | + api_schema.get_custom_operation_service(initial_method) |
| 1734 | + |
| 1735 | + |
| 1736 | +def test_extended_operations_lro_operation_service_no_polling_method(): |
| 1737 | + file_protos = generate_basic_extended_operations_setup() |
| 1738 | + |
| 1739 | + api_schema = api.API.build(file_protos) |
| 1740 | + initial_method = api_schema.services["google.extended_operations.v1.stuff.RegularService"].methods["CreateTask"] |
| 1741 | + |
| 1742 | + operation_service = api_schema.services["google.extended_operations.v1.stuff.CustomOperations"] |
| 1743 | + del operation_service.methods["Get"].options.Extensions[ex_ops_pb2.operation_polling_method] |
| 1744 | + |
| 1745 | + with pytest.raises(ValueError): |
| 1746 | + api_schema.get_custom_operation_service(initial_method) |
0 commit comments