Enhance YandexProvider functionality and tests

- Updated YandexProvider to support folder ID and IAM token authentication.
- Refactored get_translation method to handle response parsing and error handling more robustly.
- Added new tests for YandexProvider covering various scenarios including autodetection, explicit language setting, and error handling.
- Adjusted main function to accept new parameters for folder ID and IAM token usage.

Author: terryyin

tests/test_provider.py

@@ -45,17 +45,84 @@ def raise_for_status(self):
 @mock.patch('requests.Session.post')
 def test_provider_yandex_make_request(mock_requests_post):
     class MockResponse:
-        text = '"dummyjson"'
+        text = '{"translations": [{"text": "test translation"}]}'
         def raise_for_status(self):
             return False
 
     mock_requests_post.return_value = MockResponse()
     provider = YandexProvider(to_lang='en', secret_access_key='secret', folder_id='some_id')
-    provider.get_translation('test')
+    result = provider.get_translation('test')
     assert mock_requests_post.called
+    assert result == 'test translation'
 
     args, kwargs = mock_requests_post.call_args
     assert 'Authorization' in kwargs['headers']
+    assert kwargs['headers']['Authorization'] == 'Api-Key secret'
     assert kwargs['json']['folderId'] == 'some_id'
     assert kwargs['json']['targetLanguageCode'] == 'en'
-    assert kwargs['json']['texts'] == 'test'
+    assert kwargs['json']['texts'] == ['test']  # Yandex API expects texts as array
+
+
+@mock.patch('requests.Session.post')
+def test_provider_yandex_autodetect(mock_requests_post):
+    class MockResponse:
+        text = '{"translations": [{"text": "translated text"}]}'
+        def raise_for_status(self):
+            return False
+
+    mock_requests_post.return_value = MockResponse()
+    provider = YandexProvider(to_lang='fr', secret_access_key='secret')
+    result = provider.get_translation('hello')
+    assert result == 'translated text'
+
+    args, kwargs = mock_requests_post.call_args
+    # Should not include sourceLanguageCode when autodetect
+    assert 'sourceLanguageCode' not in kwargs['json']
+
+
+@mock.patch('requests.Session.post')
+def test_provider_yandex_explicit_from_lang(mock_requests_post):
+    class MockResponse:
+        text = '{"translations": [{"text": "bonjour"}]}'
+        def raise_for_status(self):
+            return False
+
+    mock_requests_post.return_value = MockResponse()
+    provider = YandexProvider(to_lang='fr', from_lang='en', secret_access_key='secret')
+    result = provider.get_translation('hello')
+    assert result == 'bonjour'
+
+    args, kwargs = mock_requests_post.call_args
+    assert kwargs['json']['sourceLanguageCode'] == 'en'
+
+
+@mock.patch('requests.Session.post')
+def test_provider_yandex_iam_token(mock_requests_post):
+    class MockResponse:
+        text = '{"translations": [{"text": "test"}]}'
+        def raise_for_status(self):
+            return False
+
+    mock_requests_post.return_value = MockResponse()
+    provider = YandexProvider(to_lang='en', secret_access_key='iam_token', is_iam=True)
+    provider.get_translation('test')
+
+    args, kwargs = mock_requests_post.call_args
+    assert kwargs['headers']['Authorization'] == 'Bearer iam_token'
+
+
+@mock.patch('requests.Session.post')
+def test_provider_yandex_error_handling(mock_requests_post):
+    import requests
+    class MockResponse:
+        text = '{"error": {"message": "Invalid API key"}}'
+        def raise_for_status(self):
+            raise requests.HTTPError("401 Client Error")
+
+    mock_requests_post.return_value = MockResponse()
+    provider = YandexProvider(to_lang='en', secret_access_key='invalid_key')
+    
+    from translate.exceptions import TranslationError
+    import pytest
+    with pytest.raises(TranslationError):
+        provider.get_translation('test')

translate/main.py

@@ -9,9 +9,8 @@
 import locale
 import sys
 
-sys.path.append(os.path.dirname(os.path.dirname(__file__)))
-from translate import Translator
-from translate.version import __version__
+from .translate import Translator
+from .version import __version__
 from .constants import CONFIG_FILE_PATH, DEFAULT_PROVIDER, TRANSLATION_FROM_DEFAULT
 
 
@@ -20,7 +19,7 @@
 
 def get_config_info(option):
     config_file_path = os.path.expanduser(CONFIG_FILE_PATH)
-    options = ('from_lang', 'to_lang', 'provider', 'secret_access_key')
+    options = ('from_lang', 'to_lang', 'provider', 'secret_access_key', 'region', 'folder_id')
     if not os.path.exists(config_file_path) or option not in options:
         return ''
 
@@ -147,8 +146,21 @@ def config_file(ctx, from_lang, to_lang, provider, secret_access_key):
     help="Region for apis",
     required=False,
 )
+@click.option(
+    'folder_id', '--folder_id',
+    default=get_config_info('folder_id'),
+    help="Folder ID for Yandex API",
+    required=False,
+)
+@click.option(
+    'is_iam', '--is_iam',
+    default=False,
+    is_flag=True,
+    help="Use IAM token authentication for Yandex API",
+    required=False,
+)
 @click.argument('text', nargs=-1, type=click.STRING, required=True)
-def main(from_lang, to_lang, provider, secret_access_key, output_only, pro, text, region):
+def main(from_lang, to_lang, provider, secret_access_key, output_only, pro, text, region, folder_id, is_iam):
     """
     Python command line tool to make online translations
 
@@ -172,6 +184,7 @@ def main(from_lang, to_lang, provider, secret_access_key, output_only, pro, text
         kwargs['pro'] = pro
         kwargs['region'] = region
         kwargs['folder_id'] = folder_id
+        kwargs['is_iam'] = is_iam
 
     translator = Translator(**kwargs)
     translation = translator.translate(text)

translate/providers/base.py

@@ -4,14 +4,16 @@
 from abc import ABCMeta, abstractmethod
 import requests
 
+from ..constants import TRANSLATION_FROM_DEFAULT
+
 
 class BaseProvider:
     __metaclass__ = ABCMeta
 
     name = ''
     base_url = ''
 
-    def __init__(self, to_lang, from_lang='en', secret_access_key=None, region=None, folder_id=None, **kwargs):
+    def __init__(self, to_lang, from_lang=TRANSLATION_FROM_DEFAULT, secret_access_key=None, region=None, folder_id=None, **kwargs):
         self.headers = {'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_6_8) AppleWebit/535.19'
                                       '(KHTML, like Gecko) Chrome/18.0.1025.168 Safari/535.19'}
         self.from_lang = from_lang

translate/providers/yandex.py

@@ -1,4 +1,5 @@
 import requests
+import json
 
 from .base import BaseProvider
 from ..constants import TRANSLATION_FROM_DEFAULT
@@ -12,48 +13,55 @@ class YandexProvider(BaseProvider):
     '''
     name = "Yandex"
     base_url = "https://translate.api.cloud.yandex.net/translate/v2/translate"
-    session = None
 
-    def __init__(self, to_lang, from_lang=None, secret_access_key=None, base_url=None, folder_id=None, is_iam=False, **kwargs):
-        super().__init__(to_lang)
-        self.from_lang = from_lang
-        self.base_url = base_url
-        self.api = secret_access_key
-        self.folder_id = folder_id # Folders used by Yandex API
-        self.is_iam = is_iam # Yandex can authorise us using API or IAM tokens
+    def __init__(self, to_lang, from_lang=TRANSLATION_FROM_DEFAULT, secret_access_key=None, base_url=None, folder_id=None, is_iam=False, **kwargs):
+        super().__init__(to_lang, from_lang=from_lang, secret_access_key=secret_access_key, folder_id=folder_id, **kwargs)
+        if base_url:
+            self.base_url = base_url
+        self.folder_id = folder_id
+        self.is_iam = is_iam
 
-    def get_translation(self, text):
-        is_autodetect = False
-        if self.from_lang in ('autodetect', None):
-            self.from_lang = None 
-            is_autodetect = True # We can send Yandex nothing if we want it to detect language automatically
+    def _make_request(self, text):
+        body = {
+            "targetLanguageCode": self.to_lang,
+            "texts": [text],  # Yandex API expects texts as an array
+        }
 
-        try:
-            body = {
-                "targetLanguageCode": self.to_lang,
-                "texts": text,
-                "folderId": self.folder_id,
-            }
+        if self.folder_id:
+            body["folderId"] = self.folder_id
 
-            if not is_autodetect:
-                body["sourceLanguageCode"] = self.from_lang # Inserting source language if we're not going to autodetect it
+        if self.from_lang != TRANSLATION_FROM_DEFAULT:
+            body["sourceLanguageCode"] = self.from_lang
 
-            headers = {
-                "Content-Type": "application/json",
-            }
+        headers = {
+            "Content-Type": "application/json",
+        }
 
-            if self.is_iam:
-                headers["Authorization"] = "Bearer {0}".format(self.api) # Passing to Yandex our IAM-token
-            else:
-                headers["Authorization"] = "Api-Key {0}".format(self.api) # Passing to Yandex our API-token
+        if self.is_iam:
+            headers["Authorization"] = "Bearer {0}".format(self.secret_access_key)
+        else:
+            headers["Authorization"] = "Api-Key {0}".format(self.secret_access_key)
 
-            response = self.session.post(
-                self.base_url,
-                json=body,
-                headers=headers,
-            )
+        response = self.session.post(
+            self.base_url,
+            json=body,
+            headers=headers,
+        )
+        response.raise_for_status()
+        return json.loads(response.text)
 
-            return response.text
+    def get_translation(self, text):
+        try:
+            data = self._make_request(text)
+
+            if "translations" in data and len(data["translations"]) > 0:
+                return data["translations"][0]["text"]
+            else:
+                raise TranslationError("No translation found in response")
+        except requests.HTTPError as e:
+            raise TranslationError("HTTP error occurred: {0}".format(str(e)))
+        except (KeyError, IndexError) as e:
+            raise TranslationError("Invalid response format: {0}".format(str(e)))
         except Exception as e:
-            raise TranslationError(e)        
+            raise TranslationError(str(e))