[sdk/python] Implement getResource in the mock monitor (#5919)

CHANGELOG.md

@@ -16,6 +16,9 @@ CHANGELOG
 - Enable resource reference feature by default.
   [#5905](https://github.com/pulumi/pulumi/pull/5905)
 
+- [sdk/python] Implement getResource in the mock monitor.
+  [#5919](https://github.com/pulumi/pulumi/pull/5919)
+
 ## 2.15.4 (2020-12-08)
 
 - Fix a problem where `pulumi import` could panic on an import error due to missing error message.

sdk/python/lib/pulumi/runtime/mocks.py

@@ -18,14 +18,13 @@ Mocks for testing.
 import asyncio
 import logging
 from abc import ABC, abstractmethod
-from typing import Optional, Awaitable, Tuple, Union, Any, TYPE_CHECKING
+from typing import Dict, NamedTuple, Optional, Tuple, TYPE_CHECKING
 
-import grpc
 from google.protobuf import empty_pb2
 from . import rpc
 from .settings import Settings, configure, get_stack, get_project, get_root_resource
 from .sync_await import _sync_await
-from ..runtime.proto import engine_pb2, engine_pb2_grpc, provider_pb2, resource_pb2, resource_pb2_grpc
+from ..runtime.proto import engine_pb2, provider_pb2, resource_pb2
 from ..runtime.stack import Stack, run_pulumi_func
 
 if TYPE_CHECKING:
@@ -71,10 +70,17 @@ class Mocks(ABC):
 
 
 class MockMonitor:
+    class ResourceRegistration(NamedTuple):
+        urn: str
+        id: str
+        state: dict
+
     mocks: Mocks
+    resources: Dict[str, ResourceRegistration]
 
     def __init__(self, mocks: Mocks):
         self.mocks = mocks
+        self.resources = dict()
 
     def make_urn(self, parent: str, type_: str, name: str) -> str:
         if parent != "":
@@ -87,6 +93,14 @@ class MockMonitor:
     def Invoke(self, request):
         args = rpc.deserialize_properties(request.args)
 
+        if request.tok == "pulumi:pulumi:getResource":
+            registered_resource = self.resources.get(args["urn"])
+            if registered_resource is None:
+                raise Exception(f"unknown resource {args['urn']}")
+            ret_proto = _sync_await(rpc.serialize_properties(registered_resource._asdict(), {}))
+            fields = {"failures": None, "return": ret_proto}
+            return provider_pb2.InvokeResponse(**fields)
+
         ret = self.mocks.call(request.tok, args, request.provider)
 
         ret_proto = _sync_await(rpc.serialize_properties(ret, {}))
@@ -97,11 +111,14 @@ class MockMonitor:
     def ReadResource(self, request):
         state = rpc.deserialize_properties(request.properties)
 
-        _, state = self.mocks.new_resource(request.type, request.name, state, request.provider, request.id)
+        id_, state = self.mocks.new_resource(request.type, request.name, state, request.provider, request.id)
 
         props_proto = _sync_await(rpc.serialize_properties(state, {}))
 
         urn = self.make_urn(request.parent, request.type, request.name)
+
+        self.resources[urn] = MockMonitor.ResourceRegistration(urn, id_, state)
+
         return resource_pb2.ReadResourceResponse(urn=urn, properties=props_proto)
 
     def RegisterResource(self, request):
@@ -116,6 +133,8 @@ class MockMonitor:
 
         obj_proto = _sync_await(rpc.serialize_properties(state, {}))
 
+        self.resources[urn] = MockMonitor.ResourceRegistration(urn, id_, state)
+
         return resource_pb2.RegisterResourceResponse(urn=urn, id=id_, object=obj_proto)
 
     def RegisterResourceOutputs(self, request):

sdk/python/lib/test_with_mocks/resources.py

@@ -1,24 +1,49 @@
+from typing import Optional
+
 import pulumi
-from pulumi import Output
 
 class MyComponent(pulumi.ComponentResource):
     outprop: pulumi.Output[str]
     def __init__(self, name, inprop: pulumi.Input[str] = None, opts = None):
-        super().__init__('pkg:index:MyComponent', name, None, opts)
+        super().__init__("pkg:index:MyComponent", name, None, opts)
         if inprop is None:
-                raise TypeError("Missing required property 'inprop'")
+            raise TypeError("Missing required property 'inprop'")
         self.outprop = pulumi.Output.from_input(inprop).apply(lambda x: f"output: {x}")
 
+
 class Instance(pulumi.CustomResource):
     public_ip: pulumi.Output[str]
     def __init__(self, resource_name, name: pulumi.Input[str] = None, value: pulumi.Input[str] = None, opts = None):
-        if name is None:
-                raise TypeError("Missing required property 'name'")
+        if opts is None:
+            opts = pulumi.ResourceOptions()
+        if name is None and not opts.urn:
+            raise TypeError("Missing required property 'name'")
         __props__: dict = dict()
         __props__["public_ip"] = None
         __props__["name"] = name
         __props__["value"] = value
-        super(Instance, self).__init__('aws:ec2/instance:Instance', resource_name, __props__, opts)
+        super(Instance, self).__init__("aws:ec2/instance:Instance", resource_name, __props__, opts)
+
+
+class Module(pulumi.runtime.ResourceModule):
+    def version(self):
+        return None
+
+    def construct(self, name: str, typ: str, urn: str) -> pulumi.Resource:
+        if typ == "aws:ec2/instance:Instance":
+            return Instance(name, opts=pulumi.ResourceOptions(urn=urn))
+        else:
+            raise Exception(f"unknown resource type {typ}")
+
+
+pulumi.runtime.register_resource_module("aws", "ec2/instance", Module())
+
+
+class MyCustom(pulumi.CustomResource):
+    instance: pulumi.Output
+    def __init__(self, resource_name, props: Optional[dict] = None, opts = None):
+        super(MyCustom, self).__init__("pkg:index:MyCustom", resource_name, props, opts)
+
 
 def do_invoke():
     value = pulumi.runtime.invoke("test:index:MyFunction", props={"value": 41}).value
@@ -28,6 +53,7 @@ mycomponent = MyComponent("mycomponent", inprop="hello")
 myinstance = Instance("instance",
                       name="myvm",
                       value=pulumi.Output.secret("secret_value"))
+mycustom = MyCustom("mycustom", {"instance": myinstance})
 invoke_result = do_invoke()
 
 pulumi.export("hello", "world")

sdk/python/lib/test_with_mocks/test_testing_with_mocks.py

@@ -40,6 +40,8 @@ class MyMocks(pulumi.runtime.Mocks):
                 'public_ip': '203.0.113.12',
             }
             return ['i-1234567890abcdef0', dict(inputs, **state)]
+        elif type_ == 'pkg:index:MyCustom':
+            return [name + '_id', inputs]
         else:
             return ['', {}]
 
@@ -61,6 +63,12 @@ class TestingWithMocks(unittest.TestCase):
             self.assertEqual(ip, '203.0.113.12')
         return resources.myinstance.public_ip.apply(check_ip)
 
+    @pulumi.runtime.test
+    def test_custom_resource_reference(self):
+        def check_instance(instance):
+            self.assertIsInstance(instance, resources.Instance)
+        return resources.mycustom.instance.apply(check_instance)
+
     @pulumi.runtime.test
     def test_invoke(self):
         return self.assertEqual(resources.invoke_result, 59)