Skip to content

Commit 8b9b989

Browse files
committed
Update test_catalog_notebook.py
1 parent 0fe0d85 commit 8b9b989

File tree

1 file changed

+28
-14
lines changed

1 file changed

+28
-14
lines changed

tests/unitary/default_setup/catalog/test_catalog_notebook.py

Lines changed: 28 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,24 @@
33
# Copyright (c) 2021, 2023 Oracle and/or its affiliates.
44
# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
55

6-
from ads.catalog.notebook import NotebookCatalog, NotebookSummaryList
7-
from ads.common import auth, oci_client
8-
from ads.common.utils import random_valid_ocid
9-
from ads.config import NB_SESSION_COMPARTMENT_OCID, PROJECT_OCID
6+
import os
7+
import unittest
108
from collections import namedtuple
119
from datetime import datetime, timezone, timedelta
12-
from oci.exceptions import ServiceError
10+
from importlib import reload
1311
from unittest import mock
1412
from unittest.mock import MagicMock, Mock, patch
13+
1514
import oci
16-
import os
1715
import pandas as pd
1816
import pytest
19-
import unittest
17+
from oci.exceptions import ServiceError
18+
19+
import ads.config
20+
from ads.catalog.notebook import NotebookCatalog, NotebookSummaryList
21+
from ads.common import auth, oci_client
22+
from ads.common.utils import random_valid_ocid
23+
from ads.config import NB_SESSION_COMPARTMENT_OCID, PROJECT_OCID
2024

2125

2226
def generate_notebook_list(
@@ -78,21 +82,32 @@ class NotebookCatalogTest(unittest.TestCase):
7882

7983
nsl = NotebookSummaryList(generate_notebook_list())
8084

85+
def setUp(self) -> None:
86+
os.environ[
87+
"NB_SESSION_COMPARTMENT_OCID"
88+
] = "ocid1.compartment.oc1.<unique_ocid>"
89+
reload(ads.config)
90+
return super().setUp()
91+
92+
def tearDown(self) -> None:
93+
os.environ.pop("NB_SESSION_COMPARTMENT_OCID", None)
94+
reload(ads.config)
95+
return super().tearDown()
96+
8197
@staticmethod
82-
def generate_notebook_response_data(self, compartment_id=None, notebook_id=None):
98+
def generate_notebook_response_data(compartment_id=None, notebook_id=None):
8399
entity_item = {
84100
"compartment_id": compartment_id,
85101
"created_by": "mock_user",
86102
"defined_tags": {},
87103
"display_name": "my new notebook catalog",
88104
"freeform_tags": {},
89105
"id": notebook_id,
90-
"lifecycle_state": "",
91106
"lifecycle_state": "ACTIVE",
92107
"notebook_session_configuration_details": "",
93108
"notebook_session_url": "oci://notebook_session_url@test_namespace",
94109
"project_id": PROJECT_OCID,
95-
"time_created": self.date_time.isoformat(),
110+
"time_created": NotebookCatalogTest.date_time.isoformat(),
96111
}
97112
notebook_response = oci.data_science.models.NotebookSession(**entity_item)
98113
return notebook_response
@@ -114,7 +129,7 @@ def test_notebook_init_without_compartment_id(self, mock_client, mock_signer):
114129
def test_decorate_notebook_session_attributes(self):
115130
"""Test NotebookCatalog._decorate_notebook_session method."""
116131
notebook = self.generate_notebook_response_data(
117-
self, compartment_id=self.comp_id, notebook_id=self.notebook_id
132+
compartment_id=self.comp_id, notebook_id=self.notebook_id
118133
)
119134

120135
def generate_get_user_data(self, compartment_id=None):
@@ -175,7 +190,6 @@ def test_get_notebook_session_with_short_id(self):
175190
def mock_get_notebook_session(notebook_id=id):
176191
return Mock(
177192
data=self.generate_notebook_response_data(
178-
self,
179193
compartment_id=self.comp_id,
180194
notebook_id=short_id_index[short_id],
181195
)
@@ -300,7 +314,7 @@ def test_update_notebook_session_with_short_id(self):
300314
wrapper = namedtuple("wrapper", ["data"])
301315
client_update_notebook_session_response = wrapper(
302316
data=self.generate_notebook_response_data(
303-
self, compartment_id=self.comp_id, notebook_id=short_id_index[short_id]
317+
compartment_id=self.comp_id, notebook_id=short_id_index[short_id]
304318
)
305319
)
306320
self.notebook_catalog.ds_client.update_notebook_session = MagicMock(
@@ -347,5 +361,5 @@ def test_notebook_summary_list_filter_invalid_param(self):
347361
# selection is a notebook session instance
348362
with pytest.raises(ValueError):
349363
self.nsl.filter(
350-
selection=self.generate_notebook_response_data(self), instance=None
364+
selection=self.generate_notebook_response_data(), instance=None
351365
)

0 commit comments

Comments
 (0)