12
12
# See the License for the specific language governing permissions and
13
13
# limitations under the License.
14
14
15
+ from datetime import datetime as dt
16
+
15
17
from unittest .mock import call , MagicMock , patch
16
18
19
+ from google .cloud import storage
17
20
from google .genai import types
21
+ import pytest
18
22
19
23
import tuning_job_create
20
24
import tuning_job_get
27
31
import tuning_with_checkpoints_textgen_with_txt
28
32
29
33
34
+ GCS_OUTPUT_BUCKET = "python-docs-samples-tests"
35
+
36
+
37
+ @pytest .fixture (scope = "session" )
38
+ def output_gcs_uri () -> str :
39
+ prefix = f"text_output/{ dt .now ()} "
40
+
41
+ yield f"gs://{ GCS_OUTPUT_BUCKET } /{ prefix } "
42
+
43
+ storage_client = storage .Client ()
44
+ bucket = storage_client .get_bucket (GCS_OUTPUT_BUCKET )
45
+ blobs = bucket .list_blobs (prefix = prefix )
46
+ for blob in blobs :
47
+ blob .delete ()
48
+
49
+
30
50
@patch ("google.genai.Client" )
31
- def test_tuning_job_create (mock_genai_client : MagicMock ) -> None :
51
+ def test_tuning_job_create (mock_genai_client : MagicMock , output_gcs_uri : str ) -> None :
32
52
# Mock the API response
33
53
mock_tuning_job = types .TuningJob (
34
54
name = "test-tuning-job" ,
@@ -40,9 +60,9 @@ def test_tuning_job_create(mock_genai_client: MagicMock) -> None:
40
60
)
41
61
mock_genai_client .return_value .tunings .tune .return_value = mock_tuning_job
42
62
43
- response = tuning_job_create .create_tuning_job ()
63
+ response = tuning_job_create .create_tuning_job (output_gcs_uri = output_gcs_uri )
44
64
45
- mock_genai_client .assert_called_once_with (http_options = types .HttpOptions (api_version = "v1 " ))
65
+ mock_genai_client .assert_called_once_with (http_options = types .HttpOptions (api_version = "v1beta1 " ))
46
66
mock_genai_client .return_value .tunings .tune .assert_called_once ()
47
67
assert response == "test-tuning-job"
48
68
@@ -121,7 +141,7 @@ def test_tuning_textgen_with_txt(mock_genai_client: MagicMock) -> None:
121
141
122
142
123
143
@patch ("google.genai.Client" )
124
- def test_tuning_job_create_with_checkpoints (mock_genai_client : MagicMock ) -> None :
144
+ def test_tuning_job_create_with_checkpoints (mock_genai_client : MagicMock , output_gcs_uri : str ) -> None :
125
145
# Mock the API response
126
146
mock_tuning_job = types .TuningJob (
127
147
name = "test-tuning-job" ,
@@ -137,9 +157,9 @@ def test_tuning_job_create_with_checkpoints(mock_genai_client: MagicMock) -> Non
137
157
)
138
158
mock_genai_client .return_value .tunings .tune .return_value = mock_tuning_job
139
159
140
- response = tuning_with_checkpoints_create .create_with_checkpoints ()
160
+ response = tuning_with_checkpoints_create .create_with_checkpoints (output_gcs_uri = output_gcs_uri )
141
161
142
- mock_genai_client .assert_called_once_with (http_options = types .HttpOptions (api_version = "v1 " ))
162
+ mock_genai_client .assert_called_once_with (http_options = types .HttpOptions (api_version = "v1beta1 " ))
143
163
mock_genai_client .return_value .tunings .tune .assert_called_once ()
144
164
assert response == "test-tuning-job"
145
165
0 commit comments