add predict.py unit testing

.github/workflows/test-all.yaml

@@ -5,6 +5,14 @@ on:
   push:
 
 jobs:
+  test-predict:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+
+      - name: Test predict.py
+        run: make test_predict
+
   test-py-3-9:
     runs-on: ubuntu-latest
     steps:

.gitignore

@@ -4,6 +4,11 @@
 env/
 venv/
 
+# Python
+*.pyc
+*.pyo
+*/__pycache__/*
+
 # Pickled models
 *.model
 *.bin

Makefile

@@ -31,22 +31,26 @@ build_shell: ## Build Python 3.11 container
 	docker build --build-arg GIT_REF=${GIT_REF} -t ${NAME}_shell:${GIT_REF} -t ${NAME}_shell:latest -f shell/Dockerfile .
 
 .PHONY: test
-test: test_3_9 test_3_10 test_3_11 ## Test all container versions
+test: test_predict test_3_9 test_3_10 test_3_11 ## Test all container versions
+
+.PHONY: test_predict
+test_predict: build_shell ## Test predict script
+	docker run -i --rm -v ./tests/:/tests/ -v /tmp:/tmp ${NAME}_shell:latest python -m unittest tests.test_predict
 
 .PHONY: test_3_9
 test_3_9: build_3_9 ## Test Python 3.9 pickle
-	docker run -i --rm -v ${PWD}:${PWD} -v /tmp:/tmp ${NAME}_py_3_9:latest --model ${PWD}/tests/models/model_3_9_legacy.pkl
-	docker run -i --rm -v ${PWD}:${PWD} -v /tmp:/tmp ${NAME}_py_3_9:latest --model ${PWD}/tests/models/model_3_9.pkl
+	docker run -i --rm -v ./tests/:/tests/ -v /tmp:/tmp ${NAME}_py_3_9:latest --model /tests/models/model_3_9_legacy.pkl
+	docker run -i --rm -v ./tests/:/tests/ -v /tmp:/tmp ${NAME}_py_3_9:latest --model /tests/models/model_3_9.pkl
 
 .PHONY: test_3_10
 test_3_10: build_3_10 ## Test Python 3.10 pickle
-	docker run -i --rm -v ${PWD}:${PWD} -v /tmp:/tmp ${NAME}_py_3_10:latest --model ${PWD}/tests/models/model_3_10_legacy.pkl
-	docker run -i --rm -v ${PWD}:${PWD} -v /tmp:/tmp ${NAME}_py_3_10:latest --model ${PWD}/tests/models/model_3_10.pkl
+	docker run -i --rm -v ./tests/:/tests/ -v /tmp:/tmp ${NAME}_py_3_10:latest --model /tests/models/model_3_10_legacy.pkl
+	docker run -i --rm -v ./tests/:/tests/ -v /tmp:/tmp ${NAME}_py_3_10:latest --model /tests/models/model_3_10.pkl
 
 .PHONY: test_3_11
 test_3_11: build_3_11 ## Test Python 3.11 pickle
-	docker run -i --rm -v ${PWD}:${PWD} -v /tmp:/tmp ${NAME}_py_3_11:latest --model ${PWD}/tests/models/model_3_11_legacy.pkl
-	docker run -i --rm -v ${PWD}:${PWD} -v /tmp:/tmp ${NAME}_py_3_11:latest --model ${PWD}/tests/models/model_3_11.pkl
+	docker run -i --rm -v ./tests/:/tests/ -v /tmp:/tmp ${NAME}_py_3_11:latest --model /tests/models/model_3_11_legacy.pkl
+	docker run -i --rm -v ./tests/:/tests/ -v /tmp:/tmp ${NAME}_py_3_11:latest --model /tests/models/model_3_11.pkl
 
 .PHONY: push_latest
 push_latest: push_latest_3_9 push_latest_3_10 push_latest_3_11 ## Push latest docker containers

predict.py

@@ -10,6 +10,7 @@
 import time
 import random
 import io
+from typing import Callable
 from inspect import signature
 
 from numerapi import NumerAPI
@@ -92,7 +93,7 @@ def exit_with_help(error):
 
 
 def retry_request_with_backoff(
-    request_func: callable,
+    request_func: Callable[[], requests.Response],
     retries: int = 10,
     delay_base: float = 1.5,
     delay_exp: float = 1.5,
@@ -144,6 +145,27 @@ def get_data(dataset, output_dir):
     return live_features
 
 
+def upload_live_output(
+    predictions: pd.DataFrame,
+    post_url: str,
+    post_data: str,
+    predictions_csv_file_name: str,
+):
+    logging.info("Uploading predictions to %s", post_url)
+    csv_buffer = io.StringIO()
+    predictions.to_csv(csv_buffer)
+
+    def post_live_output():
+        csv_buffer.seek(0)
+        return requests.post(
+            post_url,
+            data=post_data,
+            files={"file": (predictions_csv_file_name, csv_buffer, "text/csv")},
+        )
+
+    retry_request_with_backoff(post_live_output)
+
+
 def main(args):
     logging.getLogger().setLevel(logging.DEBUG if args.debug else logging.INFO)
 
@@ -239,19 +261,12 @@ def main(args):
         args.output_dir, f"live_predictions-{secrets.token_hex(6)}.csv"
     )
     if args.post_url:
-        logging.info("Uploading predictions to %s", args.post_url)
-        csv_buffer = io.StringIO()
-        predictions.to_csv(csv_buffer)
-
-        def post_live_output():
-            csv_buffer.seek(0)
-            return requests.post(
-                args.post_url,
-                data=args.post_data,
-                files={"file": (predictions_csv_file_name, csv_buffer, "text/csv")},
-            )
-
-        retry_request_with_backoff(post_live_output)
+        upload_live_output(
+            predictions,
+            args.post_url,
+            args.post_data,
+            predictions_csv_file_name,
+        )
     else:
         logging.info("Saving predictions to %s", predictions_csv_file_name)
         with open(predictions_csv_file_name, "w") as f:

tests/test_predict.py

@@ -0,0 +1,48 @@
+"""
+Set up python unittests for the predict module
+"""
+
+import os
+import sys
+import unittest
+from unittest import mock
+
+import predict
+
+import pandas as pd
+
+
+class TestPredict(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def test__retry_request_with_backoff_bad_callable_return(self):
+        self.assertRaises(
+            AttributeError,
+            predict.retry_request_with_backoff,
+            lambda: "not a response object",
+        )
+
+    def test__upload_live_output(self):
+        post_mock = mock.MagicMock()
+        predict.requests.post = post_mock
+
+        post_url_str = "http://example.com"
+        data_dict = {"post": "data"}
+        filename_str = "predictions.csv"
+
+        def check_inputs(post_url, data, files):
+            assert post_url == post_url_str
+            assert data == data_dict
+            assert files == {"file": (filename_str, mock.ANY, "text/csv")}
+            response_mock = mock.MagicMock()
+            response_mock.status_code = 200
+            return response_mock
+
+        post_mock.side_effect = check_inputs
+        predict.upload_live_output(
+            pd.DataFrame({"column": [1, 2, 3]}),
+            post_url_str,
+            data_dict,
+            filename_str,
+        )