how to enable memory profiler only when testing in python?

I am using memory_profiler in python and profiling some piece of code. But I only want that to enable when testing is happening, like the function is being invoked from a test-suite. I don't want the profiler being enabled when code is running in production. Is there any way to do that?

Or I am open to generic suggestion like how to enable decorators only when the testing is happening.

Example code -

from memory_profiler import profile

@profile(precision=4)
def my_func():
    a = [1] * (10 ** 6)
    b = [2] * (2 * 10 ** 7)
    del b
    return a