Skip to content

Commit 213a728

Browse files
brandonjbjellandbrandonjbjelland
committed
checkpoint
1 parent d58e4f1 commit 213a728

7 files changed

+175
-174
lines changed

‎Makefile

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,23 @@
11
VENV_NAME=.venv
22
PYTHON=${VENV_NAME}/bin/python3
33
PIP=${VENV_NAME}/bin/pip
4+
COMMON_VARS ?= PYDEVD_DISABLE_FILE_VALIDATION=1 JUPYTER_PLATFORM_DIRS=1
45

56
SUBSTRATUS_BRANCH ?= main
7+
PROJECT_ID ?= substratus-integration-tests
68

79
.PHONY: venv
810
venv:
911
if [ ! -d "${VENV_NAME}" ]; then python3 -m venv ${VENV_NAME}; fi
1012

1113
.PHONY: install
1214
install: venv
13-
${PIP} install -r requirements.txt && \
14-
${PYTHON} -m ipykernel install --user
15+
$(COMMON_VARS) ${PIP} install -q -r requirements.txt && \
16+
$(COMMON_VARS) ${PYTHON} -m ipykernel install --user
1517

1618
.PHONY: test
1719
test: install
18-
PYDEVD_DISABLE_FILE_VALIDATION=1 JUPYTER_PLATFORM_DIRS=1 \
19-
${VENV_NAME}/bin/pytest -s --branch=$(SUBSTRATUS_BRANCH)
20+
$(COMMON_VARS) PROJECT_ID=$(PROJECT_ID) ${VENV_NAME}/bin/pytest -svvv --branch=$(SUBSTRATUS_BRANCH)
2021

2122

2223
.PHONY: freeze

‎requirements.txt

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
appnope==0.1.3
22
asttokens==2.2.1
3+
asyncio==3.4.3
34
attrs==23.1.0
45
backcall==0.2.0
56
black==23.7.0
@@ -26,6 +27,7 @@ iniconfig==2.0.0
2627
ipykernel==6.24.0
2728
ipython==8.14.0
2829
isort==5.12.0
30+
janus==1.0.0
2931
jedi==0.18.2
3032
jsonschema==4.18.4
3133
jsonschema-specifications==2023.7.1
@@ -52,6 +54,7 @@ pyasn1==0.5.0
5254
pyasn1-modules==0.3.0
5355
Pygments==2.15.1
5456
pytest==7.4.0
57+
pytest-asyncio==0.21.1
5558
pytest-dependency==0.5.1
5659
python-dateutil==2.8.2
5760
pyzmq==25.1.0
@@ -65,5 +68,7 @@ testbook==0.4.2
6568
tomli==2.0.1
6669
tornado==6.3.2
6770
traitlets==5.9.0
71+
typing_extensions==4.7.1
6872
urllib3==1.26.16
73+
watchdog==3.0.0
6974
wcwidth==0.2.6

‎scratch/capture_output_stream.py renamed to ‎tests/capture_output_stream.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ async def main():
5959
paths_to_watch = [
6060
f"{os.path.expanduser('~')}/Library/Jupyter/runtime/", # only works on mac OS
6161
"/private/var/folders/9n/1rd9yjf913s10bzn5w9mdf_m0000gn/T/", # I'm certain this is not portable as is
62+
"/tmp/",
6263
]
6364

6465
existing_files = {
@@ -78,7 +79,3 @@ async def main():
7879
queue_watcher = watch_queue(queue, existing_files)
7980

8081
await asyncio.gather(*tasks, observer_task, queue_watcher)
81-
82-
83-
# Run the main loop
84-
asyncio.run(main())

‎tests/conftest.py

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,21 @@
11
import pytest
2+
import asyncio
3+
import os
4+
import subprocess
5+
6+
from pytest_dependency import depends
7+
from testbook import testbook
8+
from google.cloud import storage
9+
import google.auth
10+
from google.auth.transport.requests import Request
11+
from testbook.testbook import TestbookNotebookClient
12+
import logging
13+
from capture_output_stream import main
14+
15+
logging.basicConfig(level=logging.INFO)
16+
# TODO(bjb): remove
17+
logging.basicConfig(level=logging.DEBUG)
18+
logger = logging.getLogger(__name__)
219

320

421
def pytest_addoption(parser):
@@ -8,3 +25,139 @@ def pytest_addoption(parser):
825
@pytest.fixture(scope="session")
926
def branch(pytestconfig):
1027
return pytestconfig.getoption("branch")
28+
29+
30+
def authenticate_to_gcp():
31+
credentials, project = google.auth.default()
32+
if not credentials.valid:
33+
if credentials.expired and credentials.refresh_token:
34+
credentials.refresh(Request())
35+
36+
if credentials.token:
37+
os.environ["GOOGLE_CREDENTIALS"] = credentials.token
38+
if os.environ.get("GOOGLE_CREDENTIALS"):
39+
credentials.token = os.environ.get("GOOGLE_CREDENTIALS")
40+
else:
41+
raise ValueError("credentials.token is empty")
42+
43+
44+
def ensure_gcp_project() -> str:
45+
project_id = subprocess.run(
46+
["gcloud", "config", "get-value", "project"],
47+
capture_output=True,
48+
text=True,
49+
).stdout.strip()
50+
51+
if not project_id or project_id == "(unset)":
52+
project_id = os.environ.get("PROJECT_ID")
53+
54+
if not project_id:
55+
raise ValueError(
56+
"Project ID is not set. Please set the PROJECT_ID environment variable."
57+
)
58+
return f"!gcloud config set project {project_id} -q"
59+
60+
61+
@pytest.fixture(scope="session")
62+
def tb_quickstart(branch):
63+
with testbook("docs/quickstart.ipynb", execute=False, timeout=1800) as tb:
64+
change_branch(tb, branch)
65+
yield tb
66+
67+
68+
@pytest.fixture(scope="session")
69+
def auth_tb_quickstart(tb_quickstart, branch):
70+
change_branch(tb_quickstart, branch)
71+
authenticate_to_gcp()
72+
tb_quickstart.inject(ensure_gcp_project())
73+
yield tb_quickstart
74+
75+
76+
@pytest.fixture(scope="module")
77+
def auth_tb_finetuning_models(branch):
78+
with testbook(
79+
"docs/walkthrough/finetuning-models.ipynb",
80+
execute=False,
81+
) as tb:
82+
change_branch(tb, branch)
83+
authenticate_to_gcp()
84+
tb.inject(ensure_gcp_project())
85+
yield tb
86+
87+
88+
@pytest.fixture(scope="module")
89+
def auth_tb_loading_datasets(branch):
90+
with testbook(
91+
"docs/walkthrough/loading-datasets.ipynb",
92+
# TODO(bjb): next iteration, try to execute=True against this nb
93+
execute=False,
94+
) as tb:
95+
change_branch(tb, branch)
96+
authenticate_to_gcp()
97+
tb.inject(ensure_gcp_project())
98+
yield tb
99+
100+
101+
@pytest.fixture(scope="module")
102+
def auth_tb_loading_models(branch):
103+
with testbook(
104+
"docs/walkthrough/loading-models.ipynb",
105+
execute=False,
106+
) as tb:
107+
change_branch(tb, branch)
108+
authenticate_to_gcp()
109+
tb.inject(ensure_gcp_project())
110+
yield tb
111+
112+
113+
@pytest.fixture(scope="module")
114+
def auth_tb_serving_models(branch):
115+
with testbook(
116+
"docs/walkthrough/serving-models.ipynb",
117+
execute=False,
118+
) as tb:
119+
change_branch(tb, branch)
120+
authenticate_to_gcp()
121+
tb.inject(ensure_gcp_project())
122+
yield tb
123+
124+
125+
@pytest.fixture(scope="session", autouse=True)
126+
def gcp_setup(auth_tb_quickstart):
127+
logger.debug("Starting gcp_setup")
128+
129+
# Create a new event loop for this fixture
130+
loop = asyncio.new_event_loop()
131+
asyncio.set_event_loop(loop)
132+
133+
logger.info("before capturing stream")
134+
main_task = loop.create_task(main()) # Create task using the created loop
135+
logger.info("output should stream")
136+
# ... rest of the code ...
137+
138+
# Close the event loop when done
139+
loop.close()
140+
141+
142+
def delete_state_lock(
143+
bucket_name="substratus-integration-tests-substratus-terraform",
144+
blob_name="primary/default.tflock",
145+
):
146+
"""Deletes the state file from the bucket."""
147+
print("deleting state lock file from bucket")
148+
storage_client = storage.Client()
149+
bucket = storage_client.bucket(bucket_name)
150+
blob = bucket.blob(blob_name)
151+
if blob.exists():
152+
blob.delete()
153+
print("Blob {} deleted.".format(blob_name))
154+
print("lock file does not exist")
155+
156+
157+
def change_branch(tb: TestbookNotebookClient, branch: str) -> None:
158+
for c in tb.cells:
159+
if "https://raw.githubusercontent.com/substratusai/substratus/main" in c.source:
160+
c.source = c.source.replace(
161+
"https://raw.githubusercontent.com/substratusai/substratus/main",
162+
f"https://raw.githubusercontent.com/substratusai/substratus/{branch}",
163+
)

‎tests/fixtures.py

Lines changed: 0 additions & 161 deletions
Original file line numberDiff line numberDiff line change
@@ -1,161 +0,0 @@
1-
import os
2-
import subprocess
3-
4-
import pytest
5-
from pytest_dependency import depends
6-
from testbook import testbook
7-
from google.cloud import storage
8-
import google.auth
9-
from google.auth.transport.requests import Request
10-
from testbook.testbook import TestbookNotebookClient
11-
12-
13-
def authenticate_to_gcp():
14-
credentials, project = google.auth.default()
15-
if not credentials.valid:
16-
if credentials.expired and credentials.refresh_token:
17-
credentials.refresh(Request())
18-
19-
if credentials.token:
20-
os.environ["GOOGLE_CREDENTIALS"] = credentials.token
21-
if os.environ.get("GOOGLE_CREDENTIALS"):
22-
credentials.token = os.environ.get("GOOGLE_CREDENTIALS")
23-
else:
24-
raise ValueError("credentials.token is empty")
25-
26-
27-
def ensure_gcp_project() -> str:
28-
project_id = subprocess.run(
29-
["gcloud", "config", "get-value", "project"],
30-
capture_output=True,
31-
text=True,
32-
).stdout.strip()
33-
34-
if project_id == "(unset)":
35-
project_id = os.environ.get("PROJECT_ID")
36-
37-
if not project_id:
38-
raise ValueError(
39-
"Project ID is not set. Please set the PROJECT_ID environment variable."
40-
)
41-
return f"!gcloud config set project {project_id} -q"
42-
43-
44-
@pytest.fixture(scope="session")
45-
def tb_quickstart(branch):
46-
with testbook("docs/quickstart.ipynb", execute=False, timeout=1800) as tb:
47-
change_branch(tb, branch)
48-
yield tb
49-
50-
51-
@pytest.fixture(scope="session")
52-
def auth_tb_quickstart(tb_quickstart, branch):
53-
change_branch(tb_quickstart, branch)
54-
authenticate_to_gcp()
55-
tb_quickstart.inject(ensure_gcp_project())
56-
yield tb_quickstart
57-
58-
59-
@pytest.fixture(scope="module")
60-
def auth_tb_finetuning_models(branch):
61-
with testbook(
62-
"docs/walkthrough/finetuning-models.ipynb",
63-
execute=False,
64-
) as tb:
65-
change_branch(tb, branch)
66-
authenticate_to_gcp()
67-
tb.inject(ensure_gcp_project())
68-
yield tb
69-
70-
71-
@pytest.fixture(scope="module")
72-
def auth_tb_loading_datasets(branch):
73-
with testbook(
74-
"docs/walkthrough/loading-datasets.ipynb",
75-
# TODO(bjb): next iteration, try to execute=True against this nb
76-
execute=False,
77-
) as tb:
78-
change_branch(tb, branch)
79-
authenticate_to_gcp()
80-
tb.inject(ensure_gcp_project())
81-
yield tb
82-
83-
84-
@pytest.fixture(scope="module")
85-
def auth_tb_loading_models(branch):
86-
with testbook(
87-
"docs/walkthrough/loading-models.ipynb",
88-
execute=False,
89-
) as tb:
90-
change_branch(tb, branch)
91-
authenticate_to_gcp()
92-
tb.inject(ensure_gcp_project())
93-
yield tb
94-
95-
96-
@pytest.fixture(scope="module")
97-
def auth_tb_serving_models(branch):
98-
with testbook(
99-
"docs/walkthrough/serving-models.ipynb",
100-
execute=False,
101-
) as tb:
102-
change_branch(tb, branch)
103-
authenticate_to_gcp()
104-
tb.inject(ensure_gcp_project())
105-
yield tb
106-
107-
108-
@pytest.fixture(scope="session", autouse=True)
109-
def gcp_setup(auth_tb_quickstart):
110-
for attempt in range(3): # Retry up to 3 times
111-
try:
112-
auth_tb_quickstart.execute_cell("installer gcp-up")
113-
# TODO(bjb): ideally we could namespace all the infra components so more than one user could be running tests at once
114-
assert "Apply complete!" in auth_tb_quickstart.cell_output_text(
115-
"installer gcp-up"
116-
)
117-
break
118-
except Exception as err:
119-
print(f"gcp-up encountered an error: {err}")
120-
if attempt == 1:
121-
delete_state_lock()
122-
continue
123-
124-
yield # teardown below the yield
125-
# NOTE(bjb): comment out the following lines to keep infra up and iterate more rapidly on tests
126-
for attempt in range(3): # Retry up to 3 times
127-
try:
128-
auth_tb_quickstart.execute_cell("installer gcp-down")
129-
assert "Apply complete!" in auth_tb_quickstart.cell_output_text(
130-
"installer gcp-down"
131-
)
132-
break
133-
except Exception as err:
134-
print(f"gcp-down encountered an error: {err}")
135-
if attempt == 1:
136-
delete_state_lock()
137-
continue
138-
139-
140-
def delete_state_lock(
141-
bucket_name="substratus-integration-tests-substratus-terraform",
142-
blob_name="primary/default.tflock",
143-
):
144-
"""Deletes the state file from the bucket."""
145-
print("deleting state lock file from bucket")
146-
storage_client = storage.Client()
147-
bucket = storage_client.bucket(bucket_name)
148-
blob = bucket.blob(blob_name)
149-
if blob.exists():
150-
blob.delete()
151-
print("Blob {} deleted.".format(blob_name))
152-
print("lock file does not exist")
153-
154-
155-
def change_branch(tb: TestbookNotebookClient, branch: str) -> None:
156-
for c in tb.cells:
157-
if "https://raw.githubusercontent.com/substratusai/substratus/main" in c.source:
158-
c.source = c.source.replace(
159-
"https://raw.githubusercontent.com/substratusai/substratus/main",
160-
f"https://raw.githubusercontent.com/substratusai/substratus/{branch}",
161-
)

0 commit comments

Comments
 (0)