1+ import pytest
2+ from google .cloud .spanner_v1 .metrics .metrics_interceptor import MetricsInterceptor
3+ from google .cloud .spanner_v1 .metrics .spanner_metrics_tracer_factory import SpannerMetricsTracerFactory
4+ from google .cloud .spanner_v1 .metrics .metrics_gfe_tracer import MetricsGfeTracer
5+ from unittest .mock import MagicMock
6+
7+
8+ @pytest .fixture
9+ def interceptor ():
10+ return MetricsInterceptor ()
11+
12+ def test_parse_resource_path_valid (interceptor ):
13+ path = "projects/my_project/instances/my_instance/databases/my_database"
14+ expected = {
15+ "project" : "my_project" ,
16+ "instance" : "my_instance" ,
17+ "database" : "my_database" ,
18+ }
19+ assert interceptor ._parse_resource_path (path ) == expected
20+
21+
22+ def test_parse_resource_path_invalid (interceptor ):
23+ path = "invalid/path"
24+ expected = {}
25+ assert interceptor ._parse_resource_path (path ) == expected
26+
27+
28+ def test_extract_resource_from_path (interceptor ):
29+ metadata = [
30+ (
31+ "google-cloud-resource-prefix" ,
32+ "projects/my_project/instances/my_instance/databases/my_database" ,
33+ )
34+ ]
35+ expected = {
36+ "project" : "my_project" ,
37+ "instance" : "my_instance" ,
38+ "database" : "my_database" ,
39+ }
40+ assert interceptor ._extract_resource_from_path (metadata ) == expected
41+
42+
43+ def test_set_metrics_tracer_attributes (interceptor ):
44+ SpannerMetricsTracerFactory .current_metrics_tracer = MockMetricTracer ()
45+ resources = {
46+ "project" : "my_project" ,
47+ "instance" : "my_instance" ,
48+ "database" : "my_database" ,
49+ }
50+
51+ interceptor ._set_metrics_tracer_attributes (resources )
52+ assert SpannerMetricsTracerFactory .current_metrics_tracer .project == "my_project"
53+ assert SpannerMetricsTracerFactory .current_metrics_tracer .instance == "my_instance"
54+ assert SpannerMetricsTracerFactory .current_metrics_tracer .database == "my_database"
55+
56+
57+ def test_intercept_without_tracer (interceptor ):
58+ mock_invoked_method = MagicMock (return_value = "response" )
59+ mock_details = MagicMock (metadata = {})
60+ response = interceptor .intercept (mock_invoked_method , "request" , mock_details )
61+
62+ assert response == "response"
63+ mock_invoked_method .assert_called_once_with ("request" , mock_details )
64+
65+
66+ def test_intercept_with_tracer (interceptor ):
67+ SpannerMetricsTracerFactory .current_metrics_tracer = MockMetricTracer ()
68+ SpannerMetricsTracerFactory .current_metrics_tracer .record_attempt_start = MagicMock ()
69+ SpannerMetricsTracerFactory .current_metrics_tracer .record_attempt_completion = MagicMock ()
70+
71+ mock_invoked_method = MagicMock (return_value = "response" )
72+ call_details = MagicMock (
73+ method = "spanner.someMethod" ,
74+ metadata = [
75+ (
76+ "google-cloud-resource-prefix" ,
77+ "projects/my_project/instances/my_instance/databases/my_database" ,
78+ )
79+ ],
80+ )
81+
82+ response = interceptor .intercept (mock_invoked_method , "request" , call_details )
83+
84+ assert response == "response"
85+ SpannerMetricsTracerFactory .current_metrics_tracer .record_attempt_start .assert_called_once ()
86+ SpannerMetricsTracerFactory .current_metrics_tracer .record_attempt_completion .assert_called_once ()
87+ mock_invoked_method .assert_called_once_with ("request" , call_details )
88+
89+ def test_intercept_with_gfe_metrics (interceptor ):
90+
91+ # Set the mock GFE tracer
92+ SpannerMetricsTracerFactory .metrics_gfe_tracer = MockGfeTracer ()
93+ SpannerMetricsTracerFactory .metrics_gfe_tracer .enabled = True
94+ mock_invoked_method = MagicMock (return_value = MagicMock (initial_metadata = lambda : [("google-cloud-resource-prefix" , "projects/my_project/instances/my_instance/databases/my_database" )]))
95+ call_details = MagicMock (
96+ method = "spanner.someMethod" ,
97+ metadata = [
98+ (
99+ "google-cloud-resource-prefix" ,
100+ "projects/my_project/instances/my_instance/databases/my_database" ,
101+ )
102+ ],
103+ )
104+
105+ response = interceptor .intercept (mock_invoked_method , "request" , call_details )
106+
107+ assert response == mock_invoked_method .return_value
108+ assert SpannerMetricsTracerFactory .metrics_gfe_tracer .metadata == [("google-cloud-resource-prefix" , "projects/my_project/instances/my_instance/databases/my_database" )]
109+ mock_invoked_method .assert_called_once_with ("request" , call_details )
110+
111+ class MockMetricTracer :
112+ def __init__ (self ):
113+ self .project = None
114+ self .instance = None
115+ self .database = None
116+ self .method = None
117+
118+ def set_project (self , project ):
119+ self .project = project
120+
121+ def set_instance (self , instance ):
122+ self .instance = instance
123+
124+ def set_database (self , database ):
125+ self .database = database
126+
127+ def set_method (self , method ):
128+ self .method = method
129+
130+ def record_attempt_start (self ):
131+ pass
132+
133+ def record_attempt_completion (self ):
134+ pass
135+ class MockGfeTracer :
136+ enabled = True
137+ metadata = []
138+
139+ def record_gfe_metrics (self , metadata ):
140+ print ("???" )
141+ self .metadata = metadata
0 commit comments