added auth and stub tests

tests/requirements.txt

@@ -0,0 +1,2 @@
+pytest==7.1.2
+pytest-xdist==2.5.0

tests/test_auth.py

@@ -0,0 +1,71 @@
+from unittest import mock
+from unittest.mock import Mock
+
+import pytest as pytest
+
+from clarifai.client.auth.helper import ClarifaiAuthHelper, clear_cache
+
+
+@pytest.fixture(autouse=True)
+def clear_caches():
+  clear_cache()
+
+
+def test_ui_default_url():
+  default = ClarifaiAuthHelper("clarifai", "main", "fake_pat")
+  assert default.ui == "https://clarifai.com"
+  assert default.pat == "fake_pat"
+
+
+@pytest.mark.parametrize(("input_url", "expected_url"), (
+    ("http://localhost:3002", "http://localhost:3002"),
+    ("https://localhost:3002", "https://localhost:3002"),
+    ("https://clarifai.com", "https://clarifai.com"),
+))
+def test_ui_urls(input_url, expected_url):
+  helper = ClarifaiAuthHelper("clarifai", "main", "fake_pat", ui=input_url)
+  assert helper.ui == expected_url
+
+
+def test_passing_no_schema_url_use_https_when_server_is_running():
+
+  def raise_exception():
+    return Mock()
+
+  with mock.patch('urllib.request.urlopen', new_callable=raise_exception):
+    helper = ClarifaiAuthHelper("clarifai", "main", "fake_pat", ui="server")
+    assert helper.ui == "https://server"
+
+
+def test_passing_no_schema_url_show_error_when_not_server_running():
+
+  def raise_exception():
+    return Mock(side_effect=Exception("http_exception"))
+
+  with mock.patch('urllib.request.urlopen', new_callable=raise_exception):
+    with pytest.raises(
+        Exception,
+        match="Could not get a valid response from url: localhost:3002, is the API running there?"
+    ):
+      ClarifaiAuthHelper("clarifai", "main", "fake_pat", ui="localhost:3002")
+
+
+def test_passing_no_schema_url_detect_http_when_SSL_in_error():
+
+  def raise_exception():
+    return Mock(side_effect=Exception("Has SSL in error"))
+
+  with mock.patch('urllib.request.urlopen', new_callable=raise_exception):
+    helper = ClarifaiAuthHelper("clarifai", "main", "fake_pat", ui="localhost:3002")
+    assert helper.ui == "http://localhost:3002"
+
+
+def test_passing_no_schema_url_require_port():
+
+  def raise_exception():
+    return Mock(side_effect=Exception("Has SSL in error"))
+
+  with mock.patch('urllib.request.urlopen', new_callable=raise_exception):
+    with pytest.raises(
+        Exception, match="When providing an insecure url it must have both host:port format"):
+      ClarifaiAuthHelper("clarifai", "main", "fake_pat", ui="localhost")

tests/test_stub.py

@@ -0,0 +1,97 @@
+from unittest import mock
+
+import grpc
+import pytest as pytest
+from clarifai_grpc.grpc.api import service_pb2
+from clarifai_grpc.grpc.api.status import status_code_pb2
+
+from clarifai.client.auth.helper import ClarifaiAuthHelper, clear_cache
+from clarifai.client.auth.stub import AuthorizedStub, RetryStub
+
+
+class MockRpcError(grpc.RpcError):
+  pass
+
+
+@pytest.fixture(autouse=True)
+def clear_caches():
+  clear_cache()
+
+
+def test_auth_unary_unary():
+  auth = ClarifaiAuthHelper("clarifai", "main", "fake_pat")
+  stub = AuthorizedStub(auth)
+  with mock.patch.object(stub.stub, 'ListInputs', spec=stub.stub.ListInputs) as mock_f:
+    req = service_pb2.ListInputsRequest()
+    req.user_app_id.app_id = 'test_auth_unary_unary'
+    stub.ListInputs(req)
+    mock_f.assert_called_with(req, metadata=auth.metadata)
+
+
+def test_auth_unary_unary_future():
+  auth = ClarifaiAuthHelper("clarifai", "main", "fake_pat")
+  stub = AuthorizedStub(auth)
+  with mock.patch.object(stub.stub, 'ListInputs', spec=stub.stub.ListInputs) as mock_f:
+    req = service_pb2.ListInputsRequest()
+    req.user_app_id.app_id = 'test_auth_unary_unary_future'
+    stub.ListInputs.future(req)
+    mock_f.future.assert_called_with(req, metadata=auth.metadata)
+
+
+def test_auth_unary_stream():
+  auth = ClarifaiAuthHelper("clarifai", "main", "fake_pat")
+  stub = AuthorizedStub(auth)
+  with mock.patch.object(stub.stub, 'StreamInputs', spec=stub.stub.StreamInputs) as mock_f:
+    req = service_pb2.StreamInputsRequest()
+    req.user_app_id.app_id = 'test_auth_unary_stream'
+    stub.StreamInputs(req)
+    mock_f.assert_called_with(req, metadata=auth.metadata)
+
+
+def test_retry_unary_unary():
+  max_attempts = 5
+  auth = ClarifaiAuthHelper("clarifai", "main", "fake_pat")
+  stub = RetryStub(AuthorizedStub(auth), max_attempts=max_attempts, backoff_time=0.0001)
+  retry_response = service_pb2.MultiInputResponse()
+  retry_response.status.code = status_code_pb2.CONN_THROTTLED
+  success_response = service_pb2.MultiInputResponse()
+  success_response.status.code = status_code_pb2.SUCCESS
+  for nfailures in range(0, max_attempts + 1):
+    mock_resps = [retry_response] * nfailures + [success_response]
+    with mock.patch.object(
+        stub.stub, 'ListInputs', spec=stub.stub.stub.ListInputs, side_effect=mock_resps) as mock_f:
+      req = service_pb2.ListInputsRequest()
+      req.user_app_id.app_id = 'test_retry_unary_unary'
+      res = stub.ListInputs(req)
+      assert mock_f.call_count == min(max_attempts, len(mock_resps))
+      if nfailures < max_attempts:
+        assert res is success_response
+      else:
+        assert res is retry_response
+
+
+def test_retry_grpcconn_unary_unary():
+  max_attempts = 5
+  auth = ClarifaiAuthHelper("clarifai", "main", "fake_pat")
+  stub = RetryStub(AuthorizedStub(auth), max_attempts=max_attempts, backoff_time=0.0001)
+  retry_response = service_pb2.MultiInputResponse()
+  retry_response.status.code = status_code_pb2.CONN_THROTTLED
+  success_response = service_pb2.MultiInputResponse()
+  success_response.status.code = status_code_pb2.SUCCESS
+  error = MockRpcError()
+  error.code = lambda: grpc.StatusCode.UNAVAILABLE
+  for nfailures in range(0, max_attempts + 1):
+    mock_resps = [error] * nfailures + [success_response]
+    with mock.patch.object(
+        stub.stub, 'ListInputs', spec=stub.stub.stub.ListInputs, side_effect=mock_resps) as mock_f:
+      req = service_pb2.ListInputsRequest()
+      req.user_app_id.app_id = 'test_retry_unary_unary'
+      try:
+        res = stub.ListInputs(req)
+      except Exception as e:
+        res = e
+      assert mock_f.call_count == min(max_attempts, len(mock_resps))
+      if nfailures < max_attempts:
+        assert res is success_response
+      else:
+        assert res is error