From 9eb1a4b6df4a703739abfbf8582ee909d900ff7e Mon Sep 17 00:00:00 2001
From: "Brandon T. Willard" <brandonwillard@users.noreply.github.com>
Date: Mon, 19 Jun 2023 12:42:09 -0500
Subject: [PATCH] Add coverage checks to CI

---
 .github/workflows/tests.yml | 53 ++++++++++++++++++++++++++++++++++++-
 pyproject.toml              | 18 +++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index aa9ca806b..aed1779de 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -31,4 +31,55 @@ jobs:
         pip install .[test]
     - name: Run tests
       run: |
-        pytest
+        pytest --cov=outlines
+    - name: Upload coverage data
+      uses: actions/upload-artifact@v3
+      with:
+        name: coverage-data
+        path: .coverage*
+        if-no-files-found: ignore
+
+  coverage:
+    name: Combine & check coverage.
+    needs: tests
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+
+      - uses: actions/setup-python@v4
+        with:
+          cache: pip
+          python-version: "3.11"
+
+      - name: Set up environment
+        run: |
+          pip install --upgrade coverage[toml]
+
+      - uses: actions/download-artifact@v3
+        with:
+          name: coverage-data
+
+      - name: Fetch master for coverage diff
+        run: |
+          git fetch --no-tags --prune origin main
+
+      - name: Combine coverage & fail if it's <100%.
+        run: |
+          # python -m coverage combine
+          python -m coverage html --skip-covered --skip-empty
+
+          # Report and write to summary.
+          python -m coverage report | sed 's/^/    /' >> $GITHUB_STEP_SUMMARY
+
+          # Report again and fail if under 100%.
+          python -m coverage report --fail-under=100
+
+      - name: Upload HTML report if check failed.
+        uses: actions/upload-artifact@v3
+        with:
+          name: html-report
+          path: htmlcov
+        if: ${{ failure() }}
diff --git a/pyproject.toml b/pyproject.toml
index 745c3e499..7eb62d473 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,6 +38,7 @@ test = [
     "diffusers",
     "pre-commit",
     "pytest",
+    "pytest-cov",
     "torch",
     "transformers"
 ]
@@ -84,3 +85,20 @@ module = [
     "transformers",
 ]
 ignore_missing_imports = true
+
+[tool.coverage.run]
+omit = [
+    "outlines/_version.py",
+    "tests/*",
+]
+branch = true
+
+[tool.coverage.report]
+omit = [
+    "tests/*",
+]
+exclude_lines = [
+    "pragma: no cover",
+    "if TYPE_CHECKING:",
+]
+show_missing = true