Fix empty cls

Author: johnbradley

src/bioclip/__main__.py

@@ -148,7 +148,7 @@ def main():
               pretrained_str=args.pretrained)
     elif args.command == 'predict':
         cls_str = args.cls
-        if os.path.exists(args.cls):
+        if args.cls and os.path.exists(args.cls):
             cls_str = create_classes_str(args.cls)
         predict(args.image_file,
                 format=args.format,

tests/test_main.py

@@ -1,7 +1,7 @@
 import unittest
 from unittest.mock import mock_open, patch
 import argparse
-from bioclip.__main__ import parse_args, Rank, create_classes_str
+from bioclip.__main__ import parse_args, Rank, create_classes_str, main
 
 
 class TestParser(unittest.TestCase):
@@ -71,3 +71,38 @@ def test_create_classes_str(self):
         data = "class1\nclass2\nclass3"
         with patch("builtins.open", mock_open(read_data=data)) as mock_file:
             self.assertEqual(create_classes_str('path/to/file'), 'class1,class2,class3')
+
+    @patch('bioclip.__main__.predict')
+    @patch('bioclip.__main__.parse_args')
+    def test_predict_no_class(self, mock_parse_args, mock_predict):
+        mock_parse_args.return_value = argparse.Namespace(command='predict', image_file='image.jpg', format='csv',
+                                                          output='stdout', rank=Rank.SPECIES, k=5, cls=None, device='cpu',
+                                                          model=None, pretrained=None)
+        main()
+        mock_predict.assert_called_with('image.jpg', format='csv', output='stdout', cls_str=None, rank=Rank.SPECIES, k=5,
+                                        device='cpu', model_str=None, pretrained_str=None)
+
+    @patch('bioclip.__main__.predict')
+    @patch('bioclip.__main__.parse_args')
+    @patch('bioclip.__main__.os')
+    def test_predict_class_list(self, mock_os, mock_parse_args, mock_predict):
+        mock_os.path.exists.return_value = False
+        mock_parse_args.return_value = argparse.Namespace(command='predict', image_file='image.jpg', format='csv',
+                                                          output='stdout', rank=Rank.SPECIES, k=5, cls='dog,fish,bird',
+                                                          device='cpu', model=None, pretrained=None)
+        main()
+        mock_predict.assert_called_with('image.jpg', format='csv', output='stdout', cls_str='dog,fish,bird', rank=Rank.SPECIES,
+                                        k=5, device='cpu', model_str=None, pretrained_str=None)
+
+    @patch('bioclip.__main__.predict')
+    @patch('bioclip.__main__.parse_args')
+    @patch('bioclip.__main__.os')
+    def test_predict_class_file(self, mock_os, mock_parse_args, mock_predict):
+        mock_os.path.exists.return_value = True
+        mock_parse_args.return_value = argparse.Namespace(command='predict', image_file='image.jpg', format='csv', 
+                                                          output='stdout', rank=Rank.SPECIES, k=5, cls='somefile.txt',
+                                                          device='cpu', model=None, pretrained=None)
+        with patch("builtins.open", mock_open(read_data='dog\nfish\nbird')) as mock_file:
+            main()
+        mock_predict.assert_called_with('image.jpg', format='csv', output='stdout', cls_str='dog,fish,bird', rank=Rank.SPECIES,
+                                        k=5, device='cpu', model_str=None, pretrained_str=None)