https://flyte.org logo
b

Brandon Segal

08/16/2022, 2:48 PM
Hey is there anyway to mock the NamedTuple output of a task and have that output then used in a seperate task? I have a situation where I have a task with some output
Copy code
class GenerateOutputs(NamedTuple):
        uri: str
    
    @reference_task(
        ...
    )
    def generate(
        input:str 
    ) -> GenerateOutputs:
    ... 
    
    @workflow
    def my_wf(
        input="default_input"
    ) -> str:
    output = generate(
        input:str=input
    )
    return output.uri
However anytime I try to mock this out with something like
Copy code
class TestNamedTuple(TestCase):
    def test_run_workflow(self):
        with task_mock(generate) as fake_generate:
            fake_generate.return_value = GenerateOutputs(uri="test_uri")
            uri = my_wf()
            self.assertIsNotNone(uri)
I get an error like :
AttributeError: 'Promise' object has no attribute 'uri'
How do I properly mock a task that has a namedTuple output with only one attribute?
k

Ketan (kumare3)

08/16/2022, 3:13 PM
Interesting usecase, I think this might need a change to support easily. Can you file a bug. A solution would be to mock it like the underlying promise object
b

Brandon Segal

08/16/2022, 3:18 PM
Yeah absolutely. To be clear if it was a named tuple of say 2 attributes we could just unpack it and there would be no issue?
k

Ketan (kumare3)

08/17/2022, 1:43 AM
I think so
12 Views