host: Add a GetConnection grpc method

Change-Id: I66ffc146e0345442379a960921421536ee29aad2

2 files changed, 22 insertions, 6 deletions

diff --git a/interact/a2dp.py b/interact/a2dp.py
index a0c1046..45c62c0 100644
--- a/interact/a2dp.py
+++ b/interact/a2dp.py
@@ -1,3 +1,5 @@
+from typing import Optional
+
 from grpc import Channel
 
 from blueberry.a2dp_grpc import A2DP
@@ -6,10 +8,14 @@ from blueberry.host_grpc import Host
 from blueberry.a2dp_pb2 import Sink, Source
 from blueberry.host_pb2 import Connection
 
-_connection: Connection = None
-_sink: Sink = None
-_source: Source = None
+_connection: Optional[Connection] = None
+_sink: Optional[Sink] = None
+_source: Optional[Source] = None
 
+def _ensure_connection(host, addr):
+    global _connection
+    if not _connection:
+        _connection = host.GetConnection(address=addr).connection
 
 def interact(channel: Channel, interaction_id: str, test: str, pts_addr: bytes):
     global _connection, _sink, _source
@@ -22,8 +28,7 @@ def interact(channel: Channel, interaction_id: str, test: str, pts_addr: bytes):
         if "SNK" in test:
             _sink = a2dp.OpenSink(connection=_connection).sink
     elif interaction_id == "TSC_AVDTP_mmi_iut_initiate_out_of_range":
-        if not _connection:
-            _connection = Connection(cookie=pts_addr)
+        _ensure_connection(host, pts_addr)
         host.Disconnect(connection=_connection)
         _connection = None
         _sink = None
diff --git a/proto/blueberry/host.proto b/proto/blueberry/host.proto
index 85697e5..61ddd8a 100644
--- a/proto/blueberry/host.proto
+++ b/proto/blueberry/host.proto
@@ -7,6 +7,7 @@ import "google/protobuf/empty.proto";
 service Host {
   rpc Reset(google.protobuf.Empty) returns (google.protobuf.Empty);
   rpc Connect(ConnectRequest) returns (ConnectResponse);
+  rpc GetConnection(GetConnectionRequest) returns (GetConnectionResponse);
   rpc Disconnect(DisconnectRequest) returns (DisconnectResponse);
   rpc ReadLocalAddress(google.protobuf.Empty) returns (ReadLocalAddressResponse);
   rpc SetConnectable(SetConnectableRequest) returns (SetConnectableResponse);
@@ -26,6 +27,16 @@ message ConnectResponse {
   }
 }
 
+message GetConnectionRequest {
+  bytes address = 1;
+}
+
+message GetConnectionResponse {
+  oneof response {
+    Connection connection = 1;
+  }
+}
+
 message DisconnectRequest {
   Connection connection = 1;
 }
@@ -40,4 +51,4 @@ message SetConnectableRequest {
   bool connectable = 1;
 }
 
-message SetConnectableResponse {}
\ No newline at end of file
+message SetConnectableResponse {}