CODEWITHSUNDEEP
pythonunit-testingflaskgenerator
lion_bash
lion_bash

Reputation: 1409

How to unit test stream generators

I am setting up unit tests for my flask app. A lot of my functions streams the output to the gui like this one:

def stream():
    def generate():
        if request.method == "POST":
            hostname = request.data.decode('utf-8')
            hostname_dn = "{}.{}".format(hostname, DOMAIN)

            logging.info("Connecting to: {}".format(hostname_dn))

            # Connect to hostname and execute create reports
            client = set_up_client()
            client.connect(hostname_dn,
                           username=USERNAME,
                           password=PASSWORD)

            cmd = ('tail -f -n0 /home/server.log')

            stdin, stdout, stderr = client.exec_command(cmd)

            for line in iter(lambda: stdout.readline(2048), ""):
                logging.info(line, end="")
                yield line
                if re.search(r'keyword', line):
                    yield 'keyword detected\n'
                    break

    return Response(stream_with_context(generate()), mimetype='text/html')

My question is how could I use assert statements to verify these functions? Since they return a stream response. Is there a way I can slap on an extra parameter in the return statement like 200 or something then use assert to verify the stream was successful?

Upvotes: 0

Views: 677

Answers (1)

Fine
Fine

Reputation: 2154

In your case, you should test how the stream func works. So I'd recommend to isolate (mock) everything that is not related to it and test how it behaves. So mock a Response object and iterate over a generator:

@patch("Response")
def test_stream(self, response_mock):
    # this should return invoked Response mock, 
    # so you need to retrieve a first argument
    res = stream()
    args, _ = res.call_args
    stream_gen = args[0]
    n_runs = 0
    for i in stream_gen:
        self.assertEqual(i, expected_value)
        n_runs += 1
    self.assertEqual(n_runs, expected_runs_count)

Upvotes: 1

Related Questions