add comment and test

Author: marcwil

memory_profiler.py

@@ -340,6 +340,7 @@ def memory_usage(proc=-1, interval=.1, timeout=None, timestamps=False,
                 if retval:
                     ret = ret, returned
                 if max_usage:
+                    # Convert the one element list produced by MemTimer to a singular value
                     ret = ret[0]
             except Exception:
                 parent = psutil.Process(os.getpid())

test/test_memory_usage.py

@@ -11,5 +11,20 @@ def test_memory_usage():
     assert ret[0] == (1, 2)
     assert ret[1] == dict(a=1)
 
+
+def test_return_value_consistency():
+    # Test return values when watching process by PID
+    pid_mem_list = memory_usage(timeout=1)
+    assert type(pid_mem_list) == list, "Memory usage of process should be a list"
+    pid_mem_max = memory_usage(timeout=1, max_usage=True)
+    assert type(pid_mem_max) == float, "Max memory usage of process should be a number"
+    # Test return values when watching callable
+    func_mem_list = memory_usage((some_func, (42,), dict(a=42)))
+    assert type(func_mem_list) == list, "Memory usage of callable should be a list"
+    func_mem_max = memory_usage((some_func, (42,), dict(a=42)), max_usage=True)
+    assert type(func_mem_max) == float, "Max memory usage of callable should be a number"
+
+
 if __name__ == "__main__":
     test_memory_usage()
+    test_return_value_consistency()