39
39
40
40
import bigquery_magics
41
41
import bigquery_magics .bigquery as magics
42
+ import bigquery_magics .graph_server as graph_server
42
43
43
44
try :
44
45
import google .cloud .bigquery_storage as bigquery_storage
@@ -677,10 +678,12 @@ def test_bigquery_graph_json_json_result(monkeypatch):
677
678
bqstorage_client_patch
678
679
), display_patch as display_mock :
679
680
run_query_mock .return_value = query_job_mock
680
- return_value = ip .run_cell_magic ("bigquery" , "--graph" , sql )
681
+ try :
682
+ return_value = ip .run_cell_magic ("bigquery" , "--graph" , sql )
683
+ finally :
684
+ graph_server .graph_server .stop_server ()
681
685
682
- # As we only support visualization with single-column queries, the visualizer should not be launched.
683
- display_mock .assert_not_called ()
686
+ display_mock .assert_called ()
684
687
685
688
assert bqstorage_mock .called # BQ storage client was used
686
689
assert isinstance (return_value , pandas .DataFrame )
@@ -729,9 +732,6 @@ def test_bigquery_graph_json_result(monkeypatch):
729
732
]
730
733
result = pandas .DataFrame (graph_json_rows , columns = ["graph_json" ])
731
734
run_query_patch = mock .patch ("bigquery_magics.bigquery._run_query" , autospec = True )
732
- graph_server_init_patch = mock .patch (
733
- "bigquery_magics.graph_server.GraphServer.init" , autospec = True
734
- )
735
735
display_patch = mock .patch ("IPython.display.display" , autospec = True )
736
736
query_job_mock = mock .create_autospec (
737
737
google .cloud .bigquery .job .QueryJob , instance = True
@@ -740,10 +740,7 @@ def test_bigquery_graph_json_result(monkeypatch):
740
740
741
741
with run_query_patch as run_query_mock , (
742
742
bqstorage_client_patch
743
- ), graph_server_init_patch as graph_server_init_mock , display_patch as display_mock :
744
- graph_server_init_mock .return_value = mock .Mock ()
745
- graph_server_init_mock .return_value .is_alive = mock .Mock ()
746
- graph_server_init_mock .return_value .is_alive .return_value = True
743
+ ), display_patch as display_mock :
747
744
run_query_mock .return_value = query_job_mock
748
745
749
746
return_value = ip .run_cell_magic ("bigquery" , "--graph" , sql )
@@ -770,7 +767,10 @@ def test_bigquery_graph_json_result(monkeypatch):
770
767
) # identifier in 3rd row of query result
771
768
772
769
# Make sure we can run a second graph query, after the graph server is already running.
773
- return_value = ip .run_cell_magic ("bigquery" , "--graph" , sql )
770
+ try :
771
+ return_value = ip .run_cell_magic ("bigquery" , "--graph" , sql )
772
+ finally :
773
+ graph_server .graph_server .stop_server ()
774
774
775
775
# Sanity check that the HTML content looks like graph visualization. Minimal check
776
776
# to allow Spanner to change its implementation without breaking this test.
@@ -841,9 +841,6 @@ def test_bigquery_graph_colab(monkeypatch):
841
841
]
842
842
result = pandas .DataFrame (graph_json_rows , columns = ["graph_json" ])
843
843
run_query_patch = mock .patch ("bigquery_magics.bigquery._run_query" , autospec = True )
844
- graph_server_init_patch = mock .patch (
845
- "bigquery_magics.graph_server.GraphServer.init" , autospec = True
846
- )
847
844
display_patch = mock .patch ("IPython.display.display" , autospec = True )
848
845
query_job_mock = mock .create_autospec (
849
846
google .cloud .bigquery .job .QueryJob , instance = True
@@ -852,10 +849,12 @@ def test_bigquery_graph_colab(monkeypatch):
852
849
853
850
with run_query_patch as run_query_mock , (
854
851
bqstorage_client_patch
855
- ), graph_server_init_patch as graph_server_init_mock , display_patch as display_mock :
852
+ ), display_patch as display_mock :
856
853
run_query_mock .return_value = query_job_mock
857
- graph_server_init_mock .return_value = None
858
- return_value = ip .run_cell_magic ("bigquery" , "--graph" , sql )
854
+ try :
855
+ return_value = ip .run_cell_magic ("bigquery" , "--graph" , sql )
856
+ finally :
857
+ graph_server .graph_server .stop_server ()
859
858
860
859
assert len (display_mock .call_args_list ) == 1
861
860
assert len (display_mock .call_args_list [0 ]) == 2
@@ -880,7 +879,6 @@ def test_bigquery_graph_colab(monkeypatch):
880
879
881
880
# Make sure we actually used colab path, not GraphServer path.
882
881
assert sys .modules ["google.colab" ].output .register_callback .called
883
- assert not graph_server_init_mock .called
884
882
885
883
assert bqstorage_mock .called # BQ storage client was used
886
884
assert isinstance (return_value , pandas .DataFrame )
@@ -902,7 +900,6 @@ def test_colab_callback():
902
900
"edges" : [],
903
901
"nodes" : [],
904
902
"query_result" : {"result" : []},
905
- "rows" : [],
906
903
"schema" : None ,
907
904
}
908
905
}
@@ -937,9 +934,6 @@ def test_bigquery_graph_missing_spanner_deps(monkeypatch):
937
934
sql = "SELECT graph_json FROM t"
938
935
result = pandas .DataFrame ([], columns = ["graph_json" ])
939
936
run_query_patch = mock .patch ("bigquery_magics.bigquery._run_query" , autospec = True )
940
- graph_server_init_patch = mock .patch (
941
- "bigquery_magics.graph_server.GraphServer.init" , autospec = True
942
- )
943
937
display_patch = mock .patch ("IPython.display.display" , autospec = True )
944
938
query_job_mock = mock .create_autospec (
945
939
google .cloud .bigquery .job .QueryJob , instance = True
@@ -948,11 +942,13 @@ def test_bigquery_graph_missing_spanner_deps(monkeypatch):
948
942
949
943
with run_query_patch as run_query_mock , (
950
944
bqstorage_client_patch
951
- ), graph_server_init_patch as graph_server_init_mock , display_patch as display_mock :
945
+ ), display_patch as display_mock :
952
946
run_query_mock .return_value = query_job_mock
953
- graph_server_init_mock .return_value = None
954
947
with pytest .raises (ImportError ):
955
- ip .run_cell_magic ("bigquery" , "--graph" , sql )
948
+ try :
949
+ ip .run_cell_magic ("bigquery" , "--graph" , sql )
950
+ finally :
951
+ graph_server .graph_server .stop_server ()
956
952
display_mock .assert_not_called ()
957
953
958
954
0 commit comments