Skip to content

Commit 318752a

Browse files
authored
test: add system tests for bigquery.ml.create_model (#2432)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery-dataframes/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> 🦕
1 parent 3012724 commit 318752a

File tree

1 file changed

+27
-0
lines changed

1 file changed

+27
-0
lines changed

tests/system/large/bigquery/test_ml.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,3 +62,30 @@ def test_generate_embedding_with_options(embedding_model):
6262
assert "ml_generate_embedding_status" in result.columns
6363
embedding = result["ml_generate_embedding_result"].to_pandas()
6464
assert len(embedding[0]) == 256
65+
66+
67+
def test_create_model_linear_regression(dataset_id):
68+
df = bpd.DataFrame({"x": [1, 2, 3], "y": [2, 4, 6]})
69+
model_name = f"{dataset_id}.linear_regression_model"
70+
71+
result = ml.create_model(
72+
model_name=model_name,
73+
options={"model_type": "LINEAR_REG", "input_label_cols": ["y"]},
74+
training_data=df,
75+
)
76+
77+
assert result["modelType"] == "LINEAR_REGRESSION"
78+
79+
80+
def test_create_model_with_transform(dataset_id):
81+
df = bpd.DataFrame({"x": [1, 2, 3], "y": [2, 4, 6]})
82+
model_name = f"{dataset_id}.transform_model"
83+
84+
result = ml.create_model(
85+
model_name=model_name,
86+
options={"model_type": "LINEAR_REG", "input_label_cols": ["y"]},
87+
training_data=df,
88+
transform=["x * 2 AS x_doubled", "y"],
89+
)
90+
91+
assert result["modelType"] == "LINEAR_REGRESSION"

0 commit comments

Comments
 (0)