Skip to content

Commit

Permalink
Merge pull request #2 from happtiq/add-unit-tests
Browse files Browse the repository at this point in the history
Add unit tests
  • Loading branch information
kaanHapptiq authored Aug 9, 2024
2 parents a99a5cd + 88f198e commit 396fd65
Show file tree
Hide file tree
Showing 15 changed files with 308 additions and 2 deletions.
7 changes: 7 additions & 0 deletions .vscode/settings.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
{
"python.testing.pytestArgs": [
"packages"
],
"python.testing.unittestEnabled": false,
"python.testing.pytestEnabled": true
}
6 changes: 6 additions & 0 deletions packages/happtiq_commons_gen_ai/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,3 +11,9 @@ dependencies = [
'google-cloud-aiplatform',
'openai'
]

[project.optional-dependencies]
dev = [
"pytest",
"pytest-mock"
]
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import vertexai
from vertexai.generative_models import GenerativeModel, Part

class VertexAiService:
class GeminiService:
def __init__(self, project_id: str, location: str, model_id: str):
vertexai.init(project=project_id, location=location)
self.model = GenerativeModel(model_name=model_id)
Expand Down
Empty file.
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
import pytest
from unittest.mock import MagicMock, patch
from happtiq_commons_gen_ai.chatgpt_api_service import ChatgptApiService

@pytest.fixture
def chatgpt_api_service():
return ChatgptApiService(api_key="test-api-key", model="some-model")

def mock_openai_client():
mock = MagicMock()
mock.chat.completions.create.return_value = {
"choices": [
{
"message": {
"content": "Generated text"
}
}
]
}
return mock

def test_send_prompt_to_api(chatgpt_api_service, monkeypatch):
monkeypatch.setattr(chatgpt_api_service, "client", mock_openai_client())

system_message = "You are a helpful assistant."
text = "Test prompt"

response = chatgpt_api_service.send_prompt_to_api(system_message, text)
assert response == "Generated text"
chatgpt_api_service.client.chat.completions.create.assert_called_once_with(
model="some-model",
messages=[{"role": "system", "content": system_message}, {"role": "user", "content": text}]
)

def test_send_prompt_to_api_error(chatgpt_api_service, monkeypatch):
mock_client = MagicMock()
mock_client.chat.completions.create.side_effect = Exception("Error generating from chatgpt")
monkeypatch.setattr(chatgpt_api_service, "client", mock_client)

system_message = "You are a helpful assistant."
text = "Test prompt"

with pytest.raises(Exception):
chatgpt_api_service.send_prompt_to_api(system_message, text)
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
import pytest
from unittest.mock import MagicMock, patch
from happtiq_commons_gen_ai.gemini_service import GeminiService

@pytest.fixture
def vertex_ai_service(monkeypatch):
monkeypatch.setattr("vertexai.init", MagicMock())
return GeminiService(project_id="some-project", location="some-location", model_id="some-model")

def mockGenerativeModel():
mock = MagicMock()
mock.generate_content.return_value = MagicMock(text="Generated text")
return mock

def test_send_single_data_prompt_to_api(vertex_ai_service, monkeypatch):
monkeypatch.setattr(vertex_ai_service, "model", mockGenerativeModel())

prompt = "Test prompt"
data_file_gcs_uri = "gs://test-bucket/test-file.txt"
mime_type = "text/plain"

response = vertex_ai_service.send_single_data_prompt_to_api(prompt, data_file_gcs_uri, mime_type)
assert response == "Generated text"
vertex_ai_service.model.generate_content.assert_called_once()

def test_send_data_prompt_to_api(vertex_ai_service, monkeypatch):
monkeypatch.setattr(vertex_ai_service, "model", mockGenerativeModel())
prompts = ["Test prompt 1", "Test prompt 2"]
data_file_gcs_uris = ["gs://test-bucket/test-file1.txt", "gs://test-bucket/test-file2.txt"]
mime_types = ["text/plain", "text/plain"]
response = vertex_ai_service.send_data_prompt_to_api(prompts, data_file_gcs_uris, mime_types)
assert response == "Generated text"
vertex_ai_service.model.generate_content.assert_called_once()
5 changes: 5 additions & 0 deletions packages/happtiq_commons_google_cloud/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -17,3 +17,8 @@ dependencies = [
'google-cloud-secret-manager'
]

[project.optional-dependencies]
dev = [
"pytest",
"pytest-mock"
]
Original file line number Diff line number Diff line change
Expand Up @@ -11,5 +11,5 @@ def auth(self, file):
credentials = Credentials.from_authorized_user_file(file)
request = google.auth.transport.requests.Request()
credentials.refresh(request)
self.logger.info(f"got scopes {credentials.scopes}")
self.logger.info("Logged in")
return credentials
Empty file.
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
import pytest
from happtiq_commons_google_cloud.auth_service import GoogleOAuth2Service
from unittest.mock import MagicMock, patch
from google.oauth2.credentials import Credentials

@pytest.fixture
def auth_service():
return GoogleOAuth2Service()

def test_auth(auth_service, monkeypatch):
mock_credentials = MagicMock(spec=Credentials)
mock_credentials.refresh.return_value = None
monkeypatch.setattr(Credentials, "from_authorized_user_file", MagicMock(return_value=mock_credentials))

file = "test_credentials.json"
credentials = auth_service.auth(file)
assert credentials == mock_credentials
Credentials.from_authorized_user_file.assert_called_once_with(file)
mock_credentials.refresh.assert_called_once()

def test_auth_error(auth_service, monkeypatch):
mock_credentials = MagicMock(spec=Credentials)
mock_credentials.refresh.side_effect = Exception("Error refreshing credentials")
monkeypatch.setattr(Credentials, "from_authorized_user_file", MagicMock(return_value=mock_credentials))

file = "test_credentials.json"
with pytest.raises(Exception):
auth_service.auth(file)
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
import pytest
from unittest.mock import MagicMock, patch
from happtiq_commons_google_cloud.gcs_api_service import GcsApiService

@pytest.fixture
def gcs_api_service():
return GcsApiService(timeout=30)

def mock_storage_client():
mock_blob = MagicMock()
mock_blob.download_as_text.return_value = "Downloaded text"
mock_bucket = MagicMock()
mock_bucket.blob.return_value = mock_blob
mock_client = MagicMock()
mock_client.bucket.return_value = mock_bucket
return mock_client

def test_download_as_text(gcs_api_service, monkeypatch):
monkeypatch.setattr(gcs_api_service, "client", mock_storage_client())

bucket_name = "test-bucket"
source_blob_name = "test-file.txt"

text = gcs_api_service.download_as_text(bucket_name, source_blob_name)
assert text == "Downloaded text"
gcs_api_service.client.bucket.assert_called_once_with(bucket_name)
gcs_api_service.client.bucket().blob.assert_called_once_with(source_blob_name)

def test_upload(gcs_api_service, monkeypatch):
mock_client = MagicMock()
mock_client.bucket.return_value = MagicMock(blob=MagicMock(upload_from_string=MagicMock()))
monkeypatch.setattr(gcs_api_service, "client", mock_client)

fileContent = b"Test file content"
bucket_name = "test-bucket"
destination = "test-file.txt"
content_type = "text/plain"

result = gcs_api_service.upload(fileContent, bucket_name, destination, content_type)
assert result == f"gs://{bucket_name}/{destination}"
gcs_api_service.client.bucket.assert_called_once_with(bucket_name)
gcs_api_service.client.bucket().blob.assert_called_once_with(destination)
gcs_api_service.client.bucket().blob(destination).upload_from_string.assert_called_once_with(fileContent, content_type=content_type, timeout=30)

def test_upload_file(gcs_api_service, monkeypatch):
mock_client = MagicMock()
mock_client.bucket.return_value = MagicMock(blob=MagicMock(upload_from_filename=MagicMock()))
monkeypatch.setattr(gcs_api_service, "client", mock_client)

file_path = "test-file.txt"
bucket_name = "test-bucket"
destination = "test-file.txt"
content_type = "text/plain"

result = gcs_api_service.upload_file(file_path, bucket_name, destination, content_type)
assert result == f"gs://{bucket_name}/{destination}"
gcs_api_service.client.bucket.assert_called_once_with(bucket_name)
gcs_api_service.client.bucket().blob.assert_called_once_with(destination)
gcs_api_service.client.bucket().blob(destination).upload_from_filename.assert_called_once_with(file_path, content_type=content_type, timeout=30)

def test_upload_file_gzipped(gcs_api_service, monkeypatch):
mock_client = MagicMock()
mock_client.bucket.return_value = MagicMock(blob=MagicMock(upload_from_filename=MagicMock()))
monkeypatch.setattr(gcs_api_service, "client", mock_client)

file_path = "test-file.txt"
bucket_name = "test-bucket"
destination = "test-file.txt"
content_type = "text/plain"

result = gcs_api_service.upload_file(file_path, bucket_name, destination, content_type, gzipped=True)
assert result == f"gs://{bucket_name}/{destination}"
gcs_api_service.client.bucket.assert_called_once_with(bucket_name)
gcs_api_service.client.bucket().blob.assert_called_once_with(destination)
gcs_api_service.client.bucket().blob(destination).upload_from_filename.assert_called_once_with(file_path, content_type=content_type, timeout=30)
gcs_api_service.client.bucket().blob(destination).content_encoding = "gzip"
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
import pytest
from happtiq_commons_google_cloud.input_parser_service import InputParserService

@pytest.fixture
def input_parser_service():
return InputParserService()

def test_parse_input_with_json(input_parser_service):
data = {
"message": {
"data": "eyJ0ZXN0IjogeyJrZXkiOiAidmFsdWUiLCAiYW5vdGhlciI6ICJpbnB1dCJ9LCAidGVzdDIiOiAic29tZXRoaW5nIn0="
}
}

message_data = input_parser_service.parse_input(data)
assert message_data == {"test": {"key": "value", "another": "input"}, "test2": "something"}

def test_parse_input_error(input_parser_service):
data = {
"message": {
"data": "invalid-base64"
}
}
with pytest.raises(Exception):
input_parser_service.parse_input(data)
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
import pytest
from happtiq_commons_google_cloud.logging_utils import is_cloud_function, setup_logging
from unittest.mock import MagicMock, patch
import os

@pytest.fixture
def empty_fixture():
return

def test_is_cloud_function_with_k_service(monkeypatch):
monkeypatch.setenv('K_SERVICE', 'test-service')
assert is_cloud_function()

def test_is_cloud_function_without_k_service(monkeypatch):
monkeypatch.delenv('K_SERVICE', raising=False)
assert not is_cloud_function()
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
import pytest
from happtiq_commons_google_cloud.pubsub_api_service import PubsubApiService
from unittest.mock import MagicMock
import json


@pytest.fixture
def pubsub_api_service():
return PubsubApiService()

def test_publish_message(pubsub_api_service, monkeypatch):
mock_publisher = MagicMock()
monkeypatch.setattr(pubsub_api_service, "publisher", mock_publisher)

topic = "projects/test-project/topics/test-topic"
message = {"key": "value"}
pubsub_api_service.publish_message(topic, message)
mock_publisher.publish.assert_called_once_with(topic, json.dumps(message).encode('utf-8'))

def test_build_topic_path(pubsub_api_service):
project_id = "test-project"
topic_id = "test-topic"
topic_path = pubsub_api_service.build_topic_path(project_id, topic_id)
assert topic_path == "projects/test-project/topics/test-topic"

def test_publish_message_error(pubsub_api_service, monkeypatch):
mock_publisher = MagicMock()
mock_publisher.publish.side_effect = Exception("Error publishing message")
monkeypatch.setattr(pubsub_api_service, "publisher", mock_publisher)
topic = "projects/test-project/topics/test-topic"
message = {"key": "value"}
with pytest.raises(Exception):
pubsub_api_service.publish_message(topic, message)
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
import pytest
from happtiq_commons_google_cloud.secret_manager_api_service import SecretManagerApiService
from unittest.mock import MagicMock, patch

@pytest.fixture
def secret_manager_api_service():
return SecretManagerApiService()

def test_read_secret(secret_manager_api_service, monkeypatch):
secret_value = "test-secret"

mock_client = MagicMock()
mock_client.access_secret_version.return_value = MagicMock(payload=MagicMock(data=bytes(secret_value, "utf-8")))
monkeypatch.setattr(secret_manager_api_service, "client", mock_client)

secret_name = "some-secret"
secret_key = secret_manager_api_service.read_secret(secret_name)
assert secret_key == secret_value

def test_build_secret_name():
project_id = "test-project"
secret_id = "test-secret"
secret_name = SecretManagerApiService.build_secret_name(project_id, secret_id)
assert secret_name == "projects/test-project/secrets/test-secret/versions/latest"

def test_read_secret_error(secret_manager_api_service, monkeypatch):
mock_client = MagicMock()
mock_client.access_secret_version.side_effect = Exception("Error accessing secret")
monkeypatch.setattr(secret_manager_api_service, "client", mock_client)

secret_name = "some-secret"
with pytest.raises(Exception):
secret_manager_api_service.read_secret(secret_name)

0 comments on commit 396fd65

Please sign in to comment.