Delete __init__.py files

test_base.py: Use with statements, assert directly after creation

Author: tdesfont

tests/__init__.py

@@ -23,12 +23,9 @@ def test_build_save_index():
     data_loader = DataLoader(params["unique_id_column"], params["feature_columns"])
     (vector_ids, vectors) = data_loader.convert_df_to_vectors(input_df)
     nearest_neighbor = NearestNeighborSearch(num_dimensions=vectors.shape[1], **params)
-    tmp = NamedTemporaryFile()
-    nearest_neighbor.build_save_index(vectors=vectors, index_path=tmp.name)
-    file_exist = os.path.isfile(tmp.name)
-    tmp.close()
-    # Test if file is created as a result
-    assert file_exist
+    with NamedTemporaryFile() as tmp:
+        nearest_neighbor.build_save_index(vectors=vectors, index_path=tmp.name)
+        assert os.path.isfile(tmp.name)
 
 
 def test_find_neighbors_df():
@@ -54,15 +51,14 @@ def test_find_neighbors_df():
     data_loader = DataLoader(params["unique_id_column"], params["feature_columns"])
     (vector_ids, vectors) = data_loader.convert_df_to_vectors(input_df)
     nearest_neighbor = NearestNeighborSearch(num_dimensions=vectors.shape[1], **params)
-    tmp = NamedTemporaryFile()
-    nearest_neighbor.build_save_index(vectors=vectors, index_path=tmp.name)
-    params = {'unique_id_column': 'images', 'feature_columns': ['prediction'], 'num_neighbors': 5}
-    nearest_neighbor = NearestNeighborSearch(**index_config)
-    nearest_neighbor.load_index(tmp.name)
-    # Find nearest neighbors in input dataset
-    df = nearest_neighbor.find_neighbors_df(input_df, **params, index_vector_ids=vector_ids)
-    actual = sorted(list(df[df['input_id'] == '34719_ostrich.jpg']['neighbor_id']))
-    expected = ['107505_ostrich.jpg', '185189_ostrich.jpg', '213657_ostrich.jpg', '229350_ostrich.jpg', '34719_ostrich.jpg']
-    tmp.close()
+    with NamedTemporaryFile() as tmp:
+        nearest_neighbor.build_save_index(vectors=vectors, index_path=tmp.name)
+        params = {'unique_id_column': 'images', 'feature_columns': ['prediction'], 'num_neighbors': 5}
+        nearest_neighbor = NearestNeighborSearch(**index_config)
+        nearest_neighbor.load_index(tmp.name)
+        # Find nearest neighbors in input dataset
+        df = nearest_neighbor.find_neighbors_df(input_df, **params, index_vector_ids=vector_ids)
+        actual = sorted(list(df[df['input_id'] == '34719_ostrich.jpg']['neighbor_id']))
+        expected = ['107505_ostrich.jpg', '185189_ostrich.jpg', '213657_ostrich.jpg', '229350_ostrich.jpg', '34719_ostrich.jpg']
     assert len(actual) == len(expected)
     assert all([actual_item == expected_item for actual_item, expected_item in zip(actual, expected)])