Skip to content

Commit

Permalink
Fix unittests tests (#535)
Browse files Browse the repository at this point in the history
* fix tokenize tests

* remove temp_dir
  • Loading branch information
MXueguang authored May 3, 2021
1 parent 4831772 commit a9d6f66
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 2 deletions.
2 changes: 1 addition & 1 deletion pyserini/tokenize_json_collection.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ def main(args):
else:
tokenizer = T5Tokenizer.from_pretrained('castorini/doc2query-t5-base-msmarco')
if (os.path.isdir(args.input)):
for i, inf in enumerate(os.listdir(args.input)):
for i, inf in enumerate(sorted(os.listdir(args.input))):
if not os.path.isdir(args.output):
os.mkdir(args.output)
outf = os.path.join(args.output, 'docs{:02d}.json'.format(i))
Expand Down
10 changes: 9 additions & 1 deletion tests/test_load_qrels.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#

import os
import shutil
import unittest

from pyserini import search
Expand All @@ -26,6 +27,9 @@ def read_file_lines(path):

class TestGetQrels(unittest.TestCase):

def setUp(self):
os.environ['PYSERINI_CACHE'] = 'temp_dir'

def test_robust04(self):
qrels_path = search.get_qrels_file('robust04')
lines = read_file_lines(qrels_path)
Expand Down Expand Up @@ -247,6 +251,10 @@ def test_trec2019_bl(self):
self.assertEqual(mid_line, "853 0 2444d88d62539b0b88dc919909cb9701 2")
self.assertEqual(last_line, "885 0 fde80cb0-b4f0-11e2-bbf2-a6f9e9d79e19 0")

def tearDown(self):
if os.path.exists('temp_dir'):
shutil.rmtree('temp_dir')


if __name__ == '__main__':
unittest.main()

0 comments on commit a9d6f66

Please sign in to comment.