From 43e12631da4a7d5ff13765cfec78acd755e4370a Mon Sep 17 00:00:00 2001
From: yuvipanda <yuvipanda@gmail.com>
Date: Thu, 30 Nov 2017 20:09:14 -0800
Subject: [PATCH] Make memory limit checking tests more solid

- Use a more precise way of triggering OOM conditions
- Make sure that the image is rebuilt each time with a cachebust
- Don't have two copies of the OOM triggering script
---
 tests/memlimit.py                       | 95 ++++++++++++++++++-------
 tests/memlimit/.gitignore               |  1 +
 tests/memlimit/dockerfile/Dockerfile    |  6 +-
 tests/memlimit/dockerfile/postBuild     | 23 +++++-
 tests/memlimit/non-dockerfile/postBuild |  6 +-
 5 files changed, 94 insertions(+), 37 deletions(-)
 create mode 100644 tests/memlimit/.gitignore
 mode change 100755 => 120000 tests/memlimit/non-dockerfile/postBuild

diff --git a/tests/memlimit.py b/tests/memlimit.py
index 9f08fc4a..617cfe72 100644
--- a/tests/memlimit.py
+++ b/tests/memlimit.py
@@ -1,36 +1,77 @@
+"""
+Test that build time memory limits are actually enforced.
+
+We give the container image at least 128M of RAM (so base things like
+apt and pip can run), and then try to allocate & use 256MB in postBuild.
+This should fail!
+"""
+import os
 import subprocess
+import time
 
-def test_memlimit_nondockerfile():
+def does_build(builddir, mem_limit, mem_allocate_mb):
+    mem_allocate_mb_file = os.path.join(builddir, 'mem_allocate_mb')
+
+    # Cache bust so we actually do a rebuild each time this is run!
+    with open(os.path.join(builddir, 'cachebust'), 'w') as cachebust:
+        cachebust.write(str(time.time()))
+    try:
+        with open(mem_allocate_mb_file, 'w') as f:
+            f.write(str(mem_allocate_mb))
+        try:
+            output = subprocess.check_output(
+                [
+                    'repo2docker',
+                    '--no-run',
+                    '--build-memory-limit', '{}M'.format(mem_limit),
+                    builddir
+                ],
+                stderr=subprocess.STDOUT,
+            ).decode()
+            print(output)
+            return True
+        except subprocess.CalledProcessError as e:
+            output = e.output.decode()
+            print(output)
+            if "The command '/bin/sh -c ./postBuild' returned a non-zero code: 137" in output:
+                return False
+            else:
+                raise
+    finally:
+        os.remove(mem_allocate_mb_file)
+
+
+
+def test_memlimit_nondockerfile_fail():
     """
     Test if memory limited builds are working for non dockerfile builds
     """
-    try:
-        subprocess.check_call([
-            'repo2docker',
-            '--no-run',
-            '--build-memory-limit', '4M',
-            'tests/memlimit/non-dockerfile'
-        ])
-        # If this doesn't throw an exception, then memory limit was
-        # not enforced!
-        assert False
-    except subprocess.CalledProcessError as e:
-        assert True
+    basedir = os.path.dirname(__file__)
+    assert not does_build(
+        os.path.join(basedir, 'memlimit/non-dockerfile'),
+        128,
+        256
+    )
+    assert does_build(
+        os.path.join(basedir, 'memlimit/non-dockerfile'),
+        512,
+        256
+    )
 
 
-def test_memlimit_dockerfile():
+def test_memlimit_dockerfile_fail():
     """
-    Test if memory limited builds are working for non dockerfile builds
+    Test if memory limited builds are working for dockerfile builds
     """
-    try:
-        subprocess.check_call([
-            'repo2docker',
-            '--no-run',
-            '--build-memory-limit', '4M',
-            'tests/memlimit/dockerfile'
-        ])
-        # If this doesn't throw an exception, then memory limit was
-        # not enforced!
-        assert False
-    except subprocess.CalledProcessError as e:
-        assert True
+    basedir = os.path.dirname(__file__)
+    assert not does_build(
+        os.path.join(basedir, 'memlimit/dockerfile'),
+        128,
+        256
+    )
+
+    assert does_build(
+        os.path.join(basedir, 'memlimit/dockerfile'),
+        512,
+        256
+    )
diff --git a/tests/memlimit/.gitignore b/tests/memlimit/.gitignore
new file mode 100644
index 00000000..1320de82
--- /dev/null
+++ b/tests/memlimit/.gitignore
@@ -0,0 +1 @@
+cachebust
diff --git a/tests/memlimit/dockerfile/Dockerfile b/tests/memlimit/dockerfile/Dockerfile
index 484911e6..62275bba 100644
--- a/tests/memlimit/dockerfile/Dockerfile
+++ b/tests/memlimit/dockerfile/Dockerfile
@@ -1,4 +1,6 @@
-FROM python:3.6
+FROM ubuntu:zesty
 
-COPY postBuild /usr/local/bin/postBuild
+RUN apt-get update && apt-get install --yes python3
+
+COPY . .
 RUN ./postBuild
diff --git a/tests/memlimit/dockerfile/postBuild b/tests/memlimit/dockerfile/postBuild
index 1426cb5b..38a2e806 100755
--- a/tests/memlimit/dockerfile/postBuild
+++ b/tests/memlimit/dockerfile/postBuild
@@ -1,5 +1,22 @@
 #!/usr/bin/env python3
+"""
+Simplest program that tries to allocate a large amount of RAM.
 
-array = []
-for i in range(1024 * 1024 * 64):
-    array.append(i)
+malloc lies on Linux by default, so we use memset to force the
+kernel to actually give us real memory.
+"""
+from ctypes import cdll, c_void_p, memset
+import os
+
+libc = cdll.LoadLibrary("libc.so.6")
+libc.malloc.restype = c_void_p
+
+with open('mem_allocate_mb') as f:
+    mem_allocate_mb = int(f.read().strip())
+
+size = 1024 * 1024 * mem_allocate_mb
+print("trying to allocate {}MB".format(mem_allocate_mb))
+
+ret = libc.malloc(size)
+
+memset(ret, 0, size)
diff --git a/tests/memlimit/non-dockerfile/postBuild b/tests/memlimit/non-dockerfile/postBuild
deleted file mode 100755
index 1426cb5b..00000000
--- a/tests/memlimit/non-dockerfile/postBuild
+++ /dev/null
@@ -1,5 +0,0 @@
-#!/usr/bin/env python3
-
-array = []
-for i in range(1024 * 1024 * 64):
-    array.append(i)
diff --git a/tests/memlimit/non-dockerfile/postBuild b/tests/memlimit/non-dockerfile/postBuild
new file mode 120000
index 00000000..874b6793
--- /dev/null
+++ b/tests/memlimit/non-dockerfile/postBuild
@@ -0,0 +1 @@
+../dockerfile/postBuild
\ No newline at end of file