This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch v2-11-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/v2-11-test by this push:
new 28aebd16dff [v2-11-test] Remove tests that fail collection
28aebd16dff is described below
commit 28aebd16dff4bd474d5a5e139d8974534a6e98aa
Author: Jarek Potiuk <[email protected]>
AuthorDate: Mon Jun 23 17:53:48 2025 +0200
[v2-11-test] Remove tests that fail collection
Those two tests fail "quarantine" tests collection and we should
remove it in 2.11 - there are no provider tests anyway
---
.../amazon/aws/log/test_s3_task_handler.py | 314 -----------
tests/providers/openai/hooks/test_openai.py | 575 ---------------------
2 files changed, 889 deletions(-)
diff --git a/tests/providers/amazon/aws/log/test_s3_task_handler.py
b/tests/providers/amazon/aws/log/test_s3_task_handler.py
deleted file mode 100644
index 93d3edecc0c..00000000000
--- a/tests/providers/amazon/aws/log/test_s3_task_handler.py
+++ /dev/null
@@ -1,314 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-from __future__ import annotations
-
-import contextlib
-import copy
-import os
-from unittest import mock
-
-import boto3
-import pytest
-from botocore.exceptions import ClientError
-from moto import mock_aws
-from tests_common.test_utils.config import conf_vars
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
-
-from airflow.models import DAG, DagRun, TaskInstance
-from airflow.providers.amazon.aws.hooks.s3 import S3Hook
-from airflow.providers.amazon.aws.log.s3_task_handler import S3TaskHandler
-from airflow.providers.standard.operators.empty import EmptyOperator
-from airflow.utils.state import State, TaskInstanceState
-from airflow.utils.timezone import datetime
-
-
[email protected](autouse=True)
-def s3mock():
- with mock_aws():
- yield
-
-
[email protected]_test
-class TestS3RemoteLogIO:
- @pytest.fixture(autouse=True)
- def setup_tests(self, create_log_template, tmp_path_factory, session):
- with conf_vars({("logging", "remote_log_conn_id"): "aws_default"}):
- self.remote_log_base = "s3://bucket/remote/log/location"
- self.remote_log_location = "s3://bucket/remote/log/location/1.log"
- self.remote_log_key = "remote/log/location/1.log"
- self.local_log_location =
str(tmp_path_factory.mktemp("local-s3-log-location"))
- create_log_template("{try_number}.log")
- self.s3_task_handler = S3TaskHandler(self.local_log_location,
self.remote_log_base)
- # Verify the hook now with the config override
- self.subject = self.s3_task_handler.io
- assert self.subject.hook is not None
-
- date = datetime(2016, 1, 1)
- self.dag = DAG("dag_for_testing_s3_task_handler", schedule=None,
start_date=date)
- task = EmptyOperator(task_id="task_for_testing_s3_log_handler",
dag=self.dag)
- if AIRFLOW_V_3_0_PLUS:
- dag_run = DagRun(
- dag_id=self.dag.dag_id,
- logical_date=date,
- run_id="test",
- run_type="manual",
- )
- else:
- dag_run = DagRun(
- dag_id=self.dag.dag_id,
- execution_date=date,
- run_id="test",
- run_type="manual",
- )
- session.add(dag_run)
- session.commit()
- session.refresh(dag_run)
-
- self.ti = TaskInstance(task=task, run_id=dag_run.run_id)
- self.ti.dag_run = dag_run
- self.ti.try_number = 1
- self.ti.state = State.RUNNING
- session.add(self.ti)
- session.commit()
-
- self.conn = boto3.client("s3")
- self.conn.create_bucket(Bucket="bucket")
- yield
-
- self.dag.clear()
-
- session.query(DagRun).delete()
- if self.s3_task_handler.handler:
- with contextlib.suppress(Exception):
- os.remove(self.s3_task_handler.handler.baseFilename)
-
- def test_hook(self):
- assert isinstance(self.subject.hook, S3Hook)
- assert self.subject.hook.transfer_config.use_threads is False
-
- def test_log_exists(self):
- self.conn.put_object(Bucket="bucket", Key=self.remote_log_key,
Body=b"")
- assert self.subject.s3_log_exists(self.remote_log_location)
-
- def test_log_exists_none(self):
- assert not self.subject.s3_log_exists(self.remote_log_location)
-
- def test_log_exists_raises(self):
- assert not self.subject.s3_log_exists("s3://nonexistentbucket/foo")
-
- def test_log_exists_no_hook(self):
- subject = S3TaskHandler(self.local_log_location,
self.remote_log_base).io
- with mock.patch.object(S3Hook, "__init__", spec=S3Hook) as mock_hook:
- mock_hook.side_effect = ConnectionError("Fake: Failed to connect")
- with pytest.raises(ConnectionError, match="Fake: Failed to
connect"):
- subject.s3_log_exists(self.remote_log_location)
-
- def test_s3_read_when_log_missing(self):
- url = "s3://bucket/foo"
- with mock.patch.object(self.subject.log, "error") as mock_error:
- result = self.subject.s3_read(url, return_error=True)
- msg = (
- f"Could not read logs from {url} with error: An error occurred
(404) when calling the "
- f"HeadObject operation: Not Found"
- )
- assert result == msg
- mock_error.assert_called_once_with(msg, exc_info=True)
-
- def test_read_raises_return_error(self):
- url = "s3://nonexistentbucket/foo"
- with mock.patch.object(self.subject.log, "error") as mock_error:
- result = self.subject.s3_read(url, return_error=True)
- msg = (
- f"Could not read logs from {url} with error: An error occurred
(NoSuchBucket) when "
- f"calling the HeadObject operation: The specified bucket does
not exist"
- )
- assert result == msg
- mock_error.assert_called_once_with(msg, exc_info=True)
-
- def test_write(self):
- with mock.patch.object(self.subject.log, "error") as mock_error:
- self.subject.write("text", self.remote_log_location)
- # We shouldn't expect any error logs in the default working case.
- mock_error.assert_not_called()
- body = boto3.resource("s3").Object("bucket",
self.remote_log_key).get()["Body"].read()
-
- assert body == b"text"
-
- def test_write_existing(self):
- self.conn.put_object(Bucket="bucket", Key=self.remote_log_key,
Body=b"previous ")
- self.subject.write("text", self.remote_log_location)
- body = boto3.resource("s3").Object("bucket",
self.remote_log_key).get()["Body"].read()
-
- assert body == b"previous \ntext"
-
- def test_write_raises(self):
- url = "s3://nonexistentbucket/foo"
- with mock.patch.object(self.subject.log, "error") as mock_error:
- self.subject.write("text", url)
- mock_error.assert_called_once_with("Could not write logs to %s",
url, exc_info=True)
-
-
[email protected]_test
-class TestS3TaskHandler:
- @conf_vars({("logging", "remote_log_conn_id"): "aws_default"})
- @pytest.fixture(autouse=True)
- def setup_tests(self, create_log_template, tmp_path_factory, session):
- self.remote_log_base = "s3://bucket/remote/log/location"
- self.remote_log_location = "s3://bucket/remote/log/location/1.log"
- self.remote_log_key = "remote/log/location/1.log"
- self.local_log_location =
str(tmp_path_factory.mktemp("local-s3-log-location"))
- create_log_template("{try_number}.log")
- self.s3_task_handler = S3TaskHandler(self.local_log_location,
self.remote_log_base)
- # Verify the hook now with the config override
- assert self.s3_task_handler.io.hook is not None
-
- date = datetime(2016, 1, 1)
- self.dag = DAG("dag_for_testing_s3_task_handler", schedule=None,
start_date=date)
- task = EmptyOperator(task_id="task_for_testing_s3_log_handler",
dag=self.dag)
- if AIRFLOW_V_3_0_PLUS:
- dag_run = DagRun(
- dag_id=self.dag.dag_id,
- logical_date=date,
- run_id="test",
- run_type="manual",
- )
- else:
- dag_run = DagRun(
- dag_id=self.dag.dag_id,
- execution_date=date,
- run_id="test",
- run_type="manual",
- )
- session.add(dag_run)
- session.commit()
- session.refresh(dag_run)
-
- self.ti = TaskInstance(task=task, run_id=dag_run.run_id)
- self.ti.dag_run = dag_run
- self.ti.try_number = 1
- self.ti.state = State.RUNNING
- session.add(self.ti)
- session.commit()
-
- self.conn = boto3.client("s3")
- self.conn.create_bucket(Bucket="bucket")
- yield
-
- self.dag.clear()
-
- session.query(DagRun).delete()
- if self.s3_task_handler.handler:
- with contextlib.suppress(Exception):
- os.remove(self.s3_task_handler.handler.baseFilename)
-
- def test_set_context_raw(self):
- self.ti.raw = True
- mock_open = mock.mock_open()
- with
mock.patch("airflow.providers.amazon.aws.log.s3_task_handler.open", mock_open):
- self.s3_task_handler.set_context(self.ti)
-
- assert not self.s3_task_handler.upload_on_close
- mock_open.assert_not_called()
-
- def test_set_context_not_raw(self):
- mock_open = mock.mock_open()
- with
mock.patch("airflow.providers.amazon.aws.log.s3_task_handler.open", mock_open):
- self.s3_task_handler.set_context(self.ti)
-
- assert self.s3_task_handler.upload_on_close
-
mock_open.assert_called_once_with(os.path.join(self.local_log_location,
"1.log"), "w")
- mock_open().write.assert_not_called()
-
- def test_read(self):
- # Test what happens when we have two log files to read
- self.conn.put_object(Bucket="bucket", Key=self.remote_log_key,
Body=b"Log line\nLine 2\n")
- self.conn.put_object(
- Bucket="bucket", Key=self.remote_log_key + ".trigger.log",
Body=b"Log line 3\nLine 4\n"
- )
- ti = copy.copy(self.ti)
- ti.state = TaskInstanceState.SUCCESS
- log, metadata = self.s3_task_handler.read(ti)
-
- expected_s3_uri = f"s3://bucket/{self.remote_log_key}"
-
- if AIRFLOW_V_3_0_PLUS:
- assert log[0].event == "::group::Log message source details"
- assert expected_s3_uri in log[0].sources
- assert log[1].event == "::endgroup::"
- assert log[2].event == "Log line"
- assert log[3].event == "Line 2"
- assert log[4].event == "Log line 3"
- assert log[5].event == "Line 4"
- assert metadata == {"end_of_log": True, "log_pos": 4}
- else:
- actual = log[0][0][-1]
- assert f"*** Found logs in s3:\n*** * {expected_s3_uri}\n" in
actual
- assert actual.endswith("Line 4")
- assert metadata == [{"end_of_log": True, "log_pos": 33}]
-
- def test_read_when_s3_log_missing(self):
- ti = copy.copy(self.ti)
- ti.state = TaskInstanceState.SUCCESS
- self.s3_task_handler._read_from_logs_server =
mock.Mock(return_value=([], []))
- log, metadata = self.s3_task_handler.read(ti)
- if AIRFLOW_V_3_0_PLUS:
- assert len(log) == 2
- assert metadata == {"end_of_log": True, "log_pos": 0}
- else:
- assert len(log) == 1
- assert len(log) == len(metadata)
- actual = log[0][0][-1]
- expected = "*** No logs found on s3 for ti=<TaskInstance:
dag_for_testing_s3_task_handler.task_for_testing_s3_log_handler test
[success]>\n"
- assert expected in actual
- assert metadata[0] == {"end_of_log": True, "log_pos": 0}
-
- def test_close(self):
- self.s3_task_handler.set_context(self.ti)
- assert self.s3_task_handler.upload_on_close
-
- self.s3_task_handler.close()
- # Should not raise
- boto3.resource("s3").Object("bucket", self.remote_log_key).get()
-
- def test_close_no_upload(self):
- self.ti.raw = True
- self.s3_task_handler.set_context(self.ti)
- assert not self.s3_task_handler.upload_on_close
- self.s3_task_handler.close()
-
- with pytest.raises(ClientError):
- boto3.resource("s3").Object("bucket", self.remote_log_key).get()
-
- @pytest.mark.parametrize(
- "delete_local_copy, expected_existence_of_local_copy",
- [(True, False), (False, True)],
- )
- def test_close_with_delete_local_logs_conf(self, delete_local_copy,
expected_existence_of_local_copy):
- with conf_vars({("logging", "delete_local_logs"):
str(delete_local_copy)}):
- handler = S3TaskHandler(self.local_log_location,
self.remote_log_base)
-
- handler.log.info("test")
- handler.set_context(self.ti)
- assert handler.upload_on_close
-
- handler.close()
- assert os.path.exists(handler.handler.baseFilename) ==
expected_existence_of_local_copy
-
- def test_filename_template_for_backward_compatibility(self):
- # filename_template arg support for running the latest provider on
airflow 2
- S3TaskHandler(self.local_log_location, self.remote_log_base,
filename_template=None)
diff --git a/tests/providers/openai/hooks/test_openai.py
b/tests/providers/openai/hooks/test_openai.py
deleted file mode 100644
index 14ff7ae2131..00000000000
--- a/tests/providers/openai/hooks/test_openai.py
+++ /dev/null
@@ -1,575 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-from __future__ import annotations
-
-import os
-from unittest.mock import patch
-
-import pytest
-
-openai = pytest.importorskip("openai")
-
-from unittest.mock import mock_open
-
-from openai.pagination import SyncCursorPage
-from openai.types import CreateEmbeddingResponse, Embedding, FileDeleted,
FileObject
-from openai.types.beta import ( # type: ignore[attr-defined]
- Assistant,
- AssistantDeleted,
- Thread,
- ThreadDeleted,
- VectorStore,
- VectorStoreDeleted,
-)
-from openai.types.beta.threads import Message, Run
-from openai.types.beta.vector_stores import VectorStoreFile,
VectorStoreFileBatch, VectorStoreFileDeleted
-from openai.types.chat import ChatCompletion
-
-from airflow.models import Connection
-from airflow.providers.openai.hooks.openai import OpenAIHook
-
-ASSISTANT_ID = "test_assistant_abc123"
-ASSISTANT_NAME = "Test Assistant"
-ASSISTANT_INSTRUCTIONS = "You are a test assistant."
-THREAD_ID = "test_thread_abc123"
-MESSAGE_ID = "test_message_abc123"
-RUN_ID = "test_run_abc123"
-MODEL = "gpt-4"
-FILE_ID = "test_file_abc123"
-FILE_NAME = "test_file.pdf"
-METADATA = {"modified": "true", "user": "abc123"}
-VECTOR_STORE_ID = "test_vs_abc123"
-VECTOR_STORE_NAME = "Test Vector Store"
-VECTOR_FILE_STORE_BATCH_ID = "test_vfsb_abc123"
-
-
[email protected]
-def mock_openai_connection():
- conn_id = "openai_conn"
- conn = Connection(
- conn_id=conn_id,
- conn_type="openai",
- )
- os.environ[f"AIRFLOW_CONN_{conn.conn_id.upper()}"] = conn.get_uri()
- return conn
-
-
[email protected]
-def mock_openai_hook(mock_openai_connection):
- with patch("airflow.providers.openai.hooks.openai.OpenAI"):
- yield OpenAIHook(conn_id=mock_openai_connection.conn_id)
-
-
[email protected]
-def mock_embeddings_response():
- return CreateEmbeddingResponse(
- data=[Embedding(embedding=[0.1, 0.2, 0.3], index=0,
object="embedding")],
- model="text-embedding-ada-002-v2",
- object="list",
- usage={"prompt_tokens": 4, "total_tokens": 4},
- )
-
-
[email protected]
-def mock_completion():
- return ChatCompletion(
- id="chatcmpl-123",
- object="chat.completion",
- created=1677652288,
- model=MODEL,
- choices=[
- {
- "index": 0,
- "message": {
- "role": "assistant",
- "content": "Hello there, how may I assist you today?",
- },
- "logprobs": None,
- "finish_reason": "stop",
- }
- ],
- )
-
-
[email protected]
-def mock_assistant():
- return Assistant(
- id=ASSISTANT_ID,
- name=ASSISTANT_NAME,
- object="assistant",
- created_at=1677652288,
- model=MODEL,
- instructions=ASSISTANT_INSTRUCTIONS,
- tools=[],
- file_ids=[],
- metadata={},
- )
-
-
[email protected]
-def mock_assistant_list(mock_assistant):
- return SyncCursorPage[Assistant](data=[mock_assistant])
-
-
[email protected]
-def mock_thread():
- return Thread(id=THREAD_ID, object="thread", created_at=1698984975,
metadata={})
-
-
[email protected]
-def mock_message():
- return Message(
- id=MESSAGE_ID,
- object="thread.message",
- created_at=1698984975,
- thread_id=THREAD_ID,
- status="completed",
- role="user",
- content=[{"type": "text", "text": {"value": "Tell me something
interesting.", "annotations": []}}],
- assistant_id=ASSISTANT_ID,
- run_id=RUN_ID,
- file_ids=[],
- metadata={},
- )
-
-
[email protected]
-def mock_message_list(mock_message):
- return SyncCursorPage[Message](data=[mock_message])
-
-
[email protected]
-def mock_run():
- return Run(
- id=RUN_ID,
- object="thread.run",
- created_at=1698107661,
- assistant_id=ASSISTANT_ID,
- parallel_tool_calls=False,
- thread_id=THREAD_ID,
- status="completed",
- started_at=1699073476,
- completed_at=1699073476,
- model=MODEL,
- instructions="You are a test assistant.",
- tools=[],
- file_ids=[],
- metadata={},
- )
-
-
[email protected]
-def mock_run_list(mock_run):
- return SyncCursorPage[Run](data=[mock_run])
-
-
[email protected]
-def mock_file():
- return FileObject(
- id=FILE_ID,
- object="file",
- bytes=120000,
- created_at=1677610602,
- filename=FILE_NAME,
- purpose="assistants",
- status="processed",
- )
-
-
[email protected]
-def mock_file_list(mock_file):
- return SyncCursorPage[FileObject](data=[mock_file])
-
-
[email protected]
-def mock_vector_store():
- return VectorStore(
- id=VECTOR_STORE_ID,
- object="vector_store",
- created_at=1698107661,
- usage_bytes=123456,
- last_active_at=1698107661,
- name=VECTOR_STORE_NAME,
- bytes=123456,
- status="completed",
- file_counts={"in_progress": 0, "completed": 100, "cancelled": 0,
"failed": 0, "total": 100},
- metadata={},
- last_used_at=1698107661,
- )
-
-
[email protected]
-def mock_vector_store_list(mock_vector_store):
- return SyncCursorPage[VectorStore](data=[mock_vector_store])
-
-
[email protected]
-def mock_vector_file_store_batch():
- return VectorStoreFileBatch(
- id=VECTOR_FILE_STORE_BATCH_ID,
- object="vector_store.files_batch",
- created_at=1699061776,
- vector_store_id=VECTOR_STORE_ID,
- status="completed",
- file_counts={
- "in_progress": 0,
- "completed": 3,
- "failed": 0,
- "cancelled": 0,
- "total": 0,
- },
- )
-
-
[email protected]
-def mock_vector_file_store_list():
- return SyncCursorPage[VectorStoreFile](
- data=[
- VectorStoreFile(
- id="test-file-abc123",
- object="vector_store.file",
- created_at=1699061776,
- usage_bytes=1234,
- vector_store_id=VECTOR_STORE_ID,
- status="completed",
- last_error=None,
- ),
- VectorStoreFile(
- id="test-file-abc456",
- object="vector_store.file",
- created_at=1699061776,
- usage_bytes=1234,
- vector_store_id=VECTOR_STORE_ID,
- status="completed",
- last_error=None,
- ),
- ]
- )
-
-
-def test_create_chat_completion(mock_openai_hook, mock_completion):
- messages = [
- {"role": "system", "content": "You are a helpful assistant."},
- {"role": "user", "content": "Hello!"},
- ]
-
- mock_openai_hook.conn.chat.completions.create.return_value =
mock_completion
- completion = mock_openai_hook.create_chat_completion(model=MODEL,
messages=messages)
- choice = completion[0]
- assert choice.message.content == "Hello there, how may I assist you today?"
-
-
-def test_create_assistant(mock_openai_hook, mock_assistant):
- mock_openai_hook.conn.beta.assistants.create.return_value = mock_assistant
- assistant = mock_openai_hook.create_assistant(
- name=ASSISTANT_NAME, model=MODEL, instructions=ASSISTANT_INSTRUCTIONS
- )
- assert assistant.name == ASSISTANT_NAME
- assert assistant.model == MODEL
- assert assistant.instructions == ASSISTANT_INSTRUCTIONS
-
-
-def test_get_assistant(mock_openai_hook, mock_assistant):
- mock_openai_hook.conn.beta.assistants.retrieve.return_value =
mock_assistant
- assistant = mock_openai_hook.get_assistant(assistant_id=ASSISTANT_ID)
- assert assistant.name == ASSISTANT_NAME
- assert assistant.model == MODEL
- assert assistant.instructions == ASSISTANT_INSTRUCTIONS
-
-
-def test_get_assistants(mock_openai_hook, mock_assistant_list):
- mock_openai_hook.conn.beta.assistants.list.return_value =
mock_assistant_list
- assistants = mock_openai_hook.get_assistants()
- assert isinstance(assistants, list)
-
-
-def test_modify_assistant(mock_openai_hook, mock_assistant):
- new_assistant_name = "New Test Assistant"
- mock_assistant.name = new_assistant_name
- mock_openai_hook.conn.beta.assistants.update.return_value = mock_assistant
- assistant = mock_openai_hook.modify_assistant(assistant_id=ASSISTANT_ID,
name=new_assistant_name)
- assert assistant.name == new_assistant_name
-
-
-def test_delete_assistant(mock_openai_hook):
- delete_response = AssistantDeleted(id=ASSISTANT_ID,
object="assistant.deleted", deleted=True)
- mock_openai_hook.conn.beta.assistants.delete.return_value = delete_response
- assistant_deleted =
mock_openai_hook.delete_assistant(assistant_id=ASSISTANT_ID)
- assert assistant_deleted.deleted
-
-
-def test_create_thread(mock_openai_hook, mock_thread):
- mock_openai_hook.conn.beta.threads.create.return_value = mock_thread
- thread = mock_openai_hook.create_thread()
- assert thread.id == THREAD_ID
-
-
-def test_modify_thread(mock_openai_hook, mock_thread):
- mock_thread.metadata = METADATA
- mock_openai_hook.conn.beta.threads.update.return_value = mock_thread
- thread = mock_openai_hook.modify_thread(thread_id=THREAD_ID,
metadata=METADATA)
- assert thread.metadata.get("modified") == "true"
- assert thread.metadata.get("user") == "abc123"
-
-
-def test_delete_thread(mock_openai_hook):
- delete_response = ThreadDeleted(id=THREAD_ID, object="thread.deleted",
deleted=True)
- mock_openai_hook.conn.beta.threads.delete.return_value = delete_response
- thread_deleted = mock_openai_hook.delete_thread(thread_id=THREAD_ID)
- assert thread_deleted.deleted
-
-
-def test_create_message(mock_openai_hook, mock_message):
- role = "user"
- content = "Tell me something interesting."
- mock_openai_hook.conn.beta.threads.messages.create.return_value =
mock_message
- message = mock_openai_hook.create_message(thread_id=THREAD_ID,
content=content, role=role)
- assert message.id == MESSAGE_ID
-
-
-def test_get_messages(mock_openai_hook, mock_message_list):
- mock_openai_hook.conn.beta.threads.messages.list.return_value =
mock_message_list
- messages = mock_openai_hook.get_messages(thread_id=THREAD_ID)
- assert isinstance(messages, list)
-
-
-def test_modify_messages(mock_openai_hook, mock_message):
- mock_message.metadata = METADATA
- mock_openai_hook.conn.beta.threads.messages.update.return_value =
mock_message
- message = mock_openai_hook.modify_message(thread_id=THREAD_ID,
message_id=MESSAGE_ID, metadata=METADATA)
- assert message.metadata.get("modified") == "true"
- assert message.metadata.get("user") == "abc123"
-
-
-def test_create_run(mock_openai_hook, mock_run):
- thread_id = THREAD_ID
- assistant_id = ASSISTANT_ID
- mock_openai_hook.conn.beta.threads.runs.create.return_value = mock_run
- run = mock_openai_hook.create_run(thread_id=thread_id,
assistant_id=assistant_id)
- assert run.id == RUN_ID
-
-
-def test_create_run_and_poll(mock_openai_hook, mock_run):
- thread_id = THREAD_ID
- assistant_id = ASSISTANT_ID
- mock_openai_hook.conn.beta.threads.runs.create_and_poll.return_value =
mock_run
- run = mock_openai_hook.create_run_and_poll(thread_id=thread_id,
assistant_id=assistant_id)
- assert run.id == RUN_ID
-
-
-def test_get_runs(mock_openai_hook, mock_run_list):
- mock_openai_hook.conn.beta.threads.runs.list.return_value = mock_run_list
- runs = mock_openai_hook.get_runs(thread_id=THREAD_ID)
- assert isinstance(runs, list)
-
-
-def test_get_run_with_run_id(mock_openai_hook, mock_run):
- mock_openai_hook.conn.beta.threads.runs.retrieve.return_value = mock_run
- run = mock_openai_hook.get_run(thread_id=THREAD_ID, run_id=RUN_ID)
- assert run.id == RUN_ID
-
-
-def test_modify_run(mock_openai_hook, mock_run):
- mock_run.metadata = METADATA
- mock_openai_hook.conn.beta.threads.runs.update.return_value = mock_run
- message = mock_openai_hook.modify_run(thread_id=THREAD_ID, run_id=RUN_ID,
metadata=METADATA)
- assert message.metadata.get("modified") == "true"
- assert message.metadata.get("user") == "abc123"
-
-
-def test_create_embeddings(mock_openai_hook, mock_embeddings_response):
- text = "Sample text"
- mock_openai_hook.conn.embeddings.create.return_value =
mock_embeddings_response
- embeddings = mock_openai_hook.create_embeddings(text)
- assert embeddings == [0.1, 0.2, 0.3]
-
-
-@patch("builtins.open", new_callable=mock_open, read_data="test-data")
-def test_upload_file(mock_file_open, mock_openai_hook, mock_file):
- mock_file.name = FILE_NAME
- mock_file.purpose = "assistants"
- mock_openai_hook.conn.files.create.return_value = mock_file
- file = mock_openai_hook.upload_file(file=mock_file_open(),
purpose="assistants")
- assert file.name == FILE_NAME
- assert file.purpose == "assistants"
-
-
-def test_get_file(mock_openai_hook, mock_file):
- mock_openai_hook.conn.files.retrieve.return_value = mock_file
- file = mock_openai_hook.get_file(file_id=FILE_ID)
- assert file.id == FILE_ID
- assert file.filename == FILE_NAME
-
-
-def test_get_files(mock_openai_hook, mock_file_list):
- mock_openai_hook.conn.files.list.return_value = mock_file_list
- files = mock_openai_hook.get_files()
- assert isinstance(files, list)
-
-
-def test_delete_file(mock_openai_hook):
- delete_response = FileDeleted(id=FILE_ID, object="file", deleted=True)
- mock_openai_hook.conn.files.delete.return_value = delete_response
- file_deleted = mock_openai_hook.delete_file(file_id=FILE_ID)
- assert file_deleted.deleted
-
-
-def test_create_vector_store(mock_openai_hook, mock_vector_store):
- mock_openai_hook.conn.beta.vector_stores.create.return_value =
mock_vector_store
- vector_store = mock_openai_hook.create_vector_store(name=VECTOR_STORE_NAME)
- assert vector_store.id == VECTOR_STORE_ID
- assert vector_store.name == VECTOR_STORE_NAME
-
-
-def test_get_vector_store(mock_openai_hook, mock_vector_store):
- mock_openai_hook.conn.beta.vector_stores.retrieve.return_value =
mock_vector_store
- vector_store =
mock_openai_hook.get_vector_store(vector_store_id=VECTOR_STORE_ID)
- assert vector_store.id == VECTOR_STORE_ID
- assert vector_store.name == VECTOR_STORE_NAME
-
-
-def test_get_vector_stores(mock_openai_hook, mock_vector_store_list):
- mock_openai_hook.conn.beta.vector_stores.list.return_value =
mock_vector_store_list
- vector_stores = mock_openai_hook.get_vector_stores()
- assert isinstance(vector_stores, list)
-
-
-def test_modify_vector_store(mock_openai_hook, mock_vector_store):
- new_vector_store_name = "New Vector Store"
- mock_vector_store.name = new_vector_store_name
- mock_openai_hook.conn.beta.vector_stores.update.return_value =
mock_vector_store
- vector_store = mock_openai_hook.modify_vector_store(
- vector_store_id=VECTOR_STORE_ID, name=new_vector_store_name
- )
- assert vector_store.name == new_vector_store_name
-
-
-def test_delete_vector_store(mock_openai_hook):
- delete_response = VectorStoreDeleted(id=VECTOR_STORE_ID,
object="vector_store.deleted", deleted=True)
- mock_openai_hook.conn.beta.vector_stores.delete.return_value =
delete_response
- vector_store_deleted =
mock_openai_hook.delete_vector_store(vector_store_id=VECTOR_STORE_ID)
- assert vector_store_deleted.deleted
-
-
-def test_upload_files_to_vector_store(mock_openai_hook,
mock_vector_file_store_batch):
- files = ["file1.txt", "file2.txt", "file3.txt"]
-
mock_openai_hook.conn.beta.vector_stores.file_batches.upload_and_poll.return_value
= (
- mock_vector_file_store_batch
- )
- vector_file_store_batch = mock_openai_hook.upload_files_to_vector_store(
- vector_store_id=VECTOR_STORE_ID, files=files
- )
- assert vector_file_store_batch.id == VECTOR_FILE_STORE_BATCH_ID
- assert vector_file_store_batch.file_counts.completed == len(files)
-
-
-def test_get_vector_store_files(mock_openai_hook, mock_vector_file_store_list):
- mock_openai_hook.conn.beta.vector_stores.files.list.return_value =
mock_vector_file_store_list
- vector_file_store_list =
mock_openai_hook.get_vector_store_files(vector_store_id=VECTOR_STORE_ID)
- assert isinstance(vector_file_store_list, list)
-
-
-def test_delete_vector_store_file(mock_openai_hook):
- delete_response = VectorStoreFileDeleted(
- id="test_file_abc123", object="vector_store.file.deleted", deleted=True
- )
- mock_openai_hook.conn.beta.vector_stores.files.delete.return_value =
delete_response
- vector_store_file_deleted = mock_openai_hook.delete_vector_store_file(
- vector_store_id=VECTOR_STORE_ID, file_id=FILE_ID
- )
- assert vector_store_file_deleted.id == FILE_ID
- assert vector_store_file_deleted.deleted
-
-
-def test_openai_hook_test_connection(mock_openai_hook):
- result, message = mock_openai_hook.test_connection()
- assert result is True
- assert message == "Connection established!"
-
-
-@patch("airflow.providers.openai.hooks.openai.OpenAI")
-def test_get_conn_with_api_key_in_extra(mock_client):
- conn_id = "api_key_in_extra"
- conn = Connection(
- conn_id=conn_id,
- conn_type="openai",
- extra={"openai_client_kwargs": {"api_key": "api_key_in_extra"}},
- )
- os.environ[f"AIRFLOW_CONN_{conn.conn_id.upper()}"] = conn.get_uri()
- hook = OpenAIHook(conn_id=conn_id)
- hook.get_conn()
- mock_client.assert_called_once_with(
- api_key="api_key_in_extra",
- base_url=None,
- )
-
-
-@patch("airflow.providers.openai.hooks.openai.OpenAI")
-def test_get_conn_with_api_key_in_password(mock_client):
- conn_id = "api_key_in_password"
- conn = Connection(
- conn_id=conn_id,
- conn_type="openai",
- password="api_key_in_password",
- )
- os.environ[f"AIRFLOW_CONN_{conn.conn_id.upper()}"] = conn.get_uri()
- hook = OpenAIHook(conn_id=conn_id)
- hook.get_conn()
- mock_client.assert_called_once_with(
- api_key="api_key_in_password",
- base_url=None,
- )
-
-
-@patch("airflow.providers.openai.hooks.openai.OpenAI")
-def test_get_conn_with_base_url_in_extra(mock_client):
- conn_id = "base_url_in_extra"
- conn = Connection(
- conn_id=conn_id,
- conn_type="openai",
- extra={"openai_client_kwargs": {"base_url": "base_url_in_extra",
"api_key": "api_key_in_extra"}},
- )
- os.environ[f"AIRFLOW_CONN_{conn.conn_id.upper()}"] = conn.get_uri()
- hook = OpenAIHook(conn_id=conn_id)
- hook.get_conn()
- mock_client.assert_called_once_with(
- api_key="api_key_in_extra",
- base_url="base_url_in_extra",
- )
-
-
-@patch("airflow.providers.openai.hooks.openai.OpenAI")
-def test_get_conn_with_openai_client_kwargs(mock_client):
- conn_id = "openai_client_kwargs"
- conn = Connection(
- conn_id=conn_id,
- conn_type="openai",
- extra={
- "openai_client_kwargs": {
- "api_key": "api_key_in_extra",
- "organization": "organization_in_extra",
- }
- },
- )
- os.environ[f"AIRFLOW_CONN_{conn.conn_id.upper()}"] = conn.get_uri()
- hook = OpenAIHook(conn_id=conn_id)
- hook.get_conn()
- mock_client.assert_called_once_with(
- api_key="api_key_in_extra",
- base_url=None,
- organization="organization_in_extra",
- )
