CODEWITHSUNDEEP
pythonpython-2.7unit-testingpython-unittest
AdeleGoldberg
AdeleGoldberg

Reputation: 1339

How to unit test a python script accepting arguments

I have the following python script which can be imported into a module or can be run standalone

# MyScrpt.py
import argparse

def DoSomething(s: str):
    return s

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Learning to unit test python')
    parser.add_argument('-myarg1', type=str, dest='myarg1', required=True, help='Help yourself')
    parser.add_argument("-myarg2", action='store_true')
    parser.add_argument("-myarg3", action='store_true')
    args = parser.parse_args()

    value = ''
    value = DoSomething(args.myarg1)
    print(value)

    if args.myarg2 == True:
        value = DoSomething('Something')
    
    print(value)

    if args.myarg3 == True:
        value = DoSomething('SomethingMore')
    
    print(value)

DoSomething is a method I want to unit test. It should return exactly the string passed into it.

Question:
How can I unit test this script? Should I simply write a script to call DoSomething and assert return value? Since it is accepting arguments, is there something more extensive in python to unit test a script?

Upvotes: 1

Views: 443

Answers (1)

AKX
AKX

Reputation: 168824

Wrap the command-line entrypoint in a function that optionally accepts argv (as ArgumentParser.parse_args() does).

import argparse


def DoSomething(s: str):
    return s


def main(argv=None):
    parser = argparse.ArgumentParser(description="Learning to unit test python")
    parser.add_argument("-myarg1", type=str, dest="myarg1", required=True, help="Help yourself")
    parser.add_argument("-myarg2", action="store_true")
    parser.add_argument("-myarg3", action="store_true")
    args = parser.parse_args(argv)

    value = ""
    value = DoSomething(args.myarg1)
    print(value)

    if args.myarg2 == True:
        value = DoSomething("Something")

    print(value)

    if args.myarg3 == True:
        value = DoSomething("SomethingMore")

    print(value)


if __name__ == "__main__":
    main()

Then you can test the main function, passing it args as you like. The below script would likely run as-is with the Py.test test runner.

import io
import contextlib
from myscript import main

def test_main():
    stdout = io.StringIO()
    with contextlib.redirect_stdout(stdout):
        main(["foo", "-myarg1", "baz", "-myarg2"])
    assert stdout.getvalue() == "... something..."

Upvotes: 2

Related Questions