From 5aa434b43944e2d93a3e5dd8c115701c28b1e0e3 Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Fri, 9 Apr 2021 07:41:23 -0500
Subject: [PATCH] Fix failing mypy tests

---
 tests/fixtures/knapsack.py            |  2 +-
 tests/solvers/test_learning_solver.py | 10 ++++++----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/tests/fixtures/knapsack.py b/tests/fixtures/knapsack.py
index 04a9e93..e7a8682 100644
--- a/tests/fixtures/knapsack.py
+++ b/tests/fixtures/knapsack.py
@@ -14,7 +14,7 @@ from tests.solvers import _is_subclass_or_instance
 
 
 def get_test_pyomo_instances() -> Tuple[List[Instance], List[Any]]:
-    instances = [
+    instances: List[Instance] = [
         KnapsackInstance(
             weights=[23.0, 26.0, 20.0, 18.0],
             prices=[505.0, 352.0, 458.0, 220.0],
diff --git a/tests/solvers/test_learning_solver.py b/tests/solvers/test_learning_solver.py
index 9596e82..9f44e8c 100644
--- a/tests/solvers/test_learning_solver.py
+++ b/tests/solvers/test_learning_solver.py
@@ -5,9 +5,11 @@
 import logging
 import os
 import tempfile
+from typing import List, cast
 
 import dill
 
+from miplearn import Instance
 from miplearn.instance.picklegz import PickleGzInstance, write_pickle_gz, read_pickle_gz
 from miplearn.solvers.gurobi import GurobiSolver
 from miplearn.solvers.learning import LearningSolver
@@ -87,7 +89,7 @@ def test_parallel_solve() -> None:
 def test_solve_fit_from_disk() -> None:
     for internal_solver in get_internal_solvers():
         # Create instances and pickle them
-        instances = []
+        instances: List[Instance] = []
         for k in range(3):
             instance = _get_knapsack_instance(internal_solver)
             with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as file:
@@ -97,7 +99,7 @@ def test_solve_fit_from_disk() -> None:
         # Test: solve
         solver = LearningSolver(solver=internal_solver)
         solver.solve(instances[0])
-        instance_loaded = read_pickle_gz(instances[0].filename)
+        instance_loaded = read_pickle_gz(cast(PickleGzInstance, instances[0]).filename)
         assert len(instance_loaded.training_data) > 0
         assert instance_loaded.features.instance is not None
         assert instance_loaded.features.variables is not None
@@ -106,7 +108,7 @@ def test_solve_fit_from_disk() -> None:
         # Test: parallel_solve
         solver.parallel_solve(instances)
         for instance in instances:
-            instance_loaded = read_pickle_gz(instance.filename)
+            instance_loaded = read_pickle_gz(cast(PickleGzInstance, instance).filename)
             assert len(instance_loaded.training_data) > 0
             assert instance_loaded.features.instance is not None
             assert instance_loaded.features.variables is not None
@@ -114,7 +116,7 @@ def test_solve_fit_from_disk() -> None:
 
         # Delete temporary files
         for instance in instances:
-            os.remove(instance.filename)
+            os.remove(cast(PickleGzInstance, instance).filename)
 
 
 def test_simulate_perfect() -> None: