From c313fa1c1d638c88272321c00d1c6efaca68e4f8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Sun, 7 Apr 2019 13:36:48 +0200
Subject: [PATCH] add TargetID to Target

This will be useful later on, for example to be able to close a target
(a page) once it gets cancelled.
---
 browser.go      | 27 ++++++++++++++++++---------
 context.go      |  2 +-
 context_test.go | 38 +++++++++++---------------------------
 target.go       |  1 +
 4 files changed, 31 insertions(+), 37 deletions(-)

diff --git a/browser.go b/browser.go
index 2080623..837bd7d 100644
--- a/browser.go
+++ b/browser.go
@@ -34,7 +34,7 @@ type Browser struct {
 	// tabQueue is the queue used to create new target handlers, once a new
 	// tab is created and attached to. The newly created Target is sent back
 	// via tabResult.
-	tabQueue  chan target.SessionID
+	tabQueue  chan newTab
 	tabResult chan *Target
 
 	// cmdQueue is the outgoing command queue.
@@ -45,6 +45,11 @@ type Browser struct {
 	errf func(string, ...interface{})
 }
 
+type newTab struct {
+	targetID  target.ID
+	sessionID target.SessionID
+}
+
 type cmdJob struct {
 	msg  *cdproto.Message
 	resp chan *cdproto.Message
@@ -60,7 +65,7 @@ func NewBrowser(ctx context.Context, urlstr string, opts ...BrowserOption) (*Bro
 	b := &Browser{
 		conn: conn,
 
-		tabQueue:  make(chan target.SessionID, 1),
+		tabQueue:  make(chan newTab, 1),
 		tabResult: make(chan *Target, 1),
 
 		cmdQueue: make(chan cmdJob),
@@ -105,11 +110,14 @@ func (b *Browser) send(method cdproto.MethodType, params easyjson.RawMessage) er
 	return b.conn.Write(msg)
 }
 
-func (b *Browser) newExecutorForTarget(ctx context.Context, sessionID target.SessionID) *Target {
+func (b *Browser) newExecutorForTarget(ctx context.Context, targetID target.ID, sessionID target.SessionID) *Target {
+	if targetID == "" {
+		panic("empty target ID")
+	}
 	if sessionID == "" {
 		panic("empty session ID")
 	}
-	b.tabQueue <- sessionID
+	b.tabQueue <- newTab{targetID, sessionID}
 	return <-b.tabResult
 }
 
@@ -231,13 +239,14 @@ func (b *Browser) run(ctx context.Context) {
 		pages := make(map[target.SessionID]*Target, 1024)
 		for {
 			select {
-			case sessionID := <-b.tabQueue:
-				if _, ok := pages[sessionID]; ok {
-					b.errf("executor for %q already exists", sessionID)
+			case tab := <-b.tabQueue:
+				if _, ok := pages[tab.sessionID]; ok {
+					b.errf("executor for %q already exists", tab.sessionID)
 				}
 				t := &Target{
 					browser:   b,
-					SessionID: sessionID,
+					TargetID:  tab.targetID,
+					SessionID: tab.sessionID,
 
 					eventQueue: make(chan *cdproto.Message, 1024),
 					waitQueue:  make(chan func(cur *cdp.Frame) bool, 1024),
@@ -247,7 +256,7 @@ func (b *Browser) run(ctx context.Context) {
 					errf: b.errf,
 				}
 				go t.run(ctx)
-				pages[sessionID] = t
+				pages[tab.sessionID] = t
 				b.tabResult <- t
 			case event := <-tabEventQueue:
 				page, ok := pages[event.sessionID]
diff --git a/context.go b/context.go
index ab10005..6a4b739 100644
--- a/context.go
+++ b/context.go
@@ -115,7 +115,7 @@ func (c *Context) newSession(ctx context.Context, first bool) error {
 	if err != nil {
 		return err
 	}
-	c.Target = c.Browser.newExecutorForTarget(ctx, sessionID)
+	c.Target = c.Browser.newExecutorForTarget(ctx, targetID, sessionID)
 
 	// enable domains
 	for _, enable := range []Action{
diff --git a/context_test.go b/context_test.go
index e970dee..8dcade2 100644
--- a/context_test.go
+++ b/context_test.go
@@ -9,47 +9,31 @@ func TestTargets(t *testing.T) {
 	t.Parallel()
 
 	// Start one browser with one tab.
-	ctx1, cancel := NewContext(context.Background())
-	defer cancel()
+	ctx1, cancel1 := NewContext(context.Background())
+	defer cancel1()
 	if err := Run(ctx1); err != nil {
 		t.Fatal(err)
 	}
 
-	{
-		infos, err := Targets(ctx1)
+	wantTargets := func(ctx context.Context, want int) {
+		infos, err := Targets(ctx)
 		if err != nil {
 			t.Fatal(err)
 		}
-		if want, got := 1, len(infos); want != got {
+		if got := len(infos); want != got {
 			t.Fatalf("want %d targets, got %d", want, got)
 		}
 	}
+	wantTargets(ctx1, 1)
 
-	// Start a second tab on the same browser
-	ctx2, cancel := NewContext(ctx1)
-	defer cancel()
+	// Start a second tab on the same browser.
+	ctx2, cancel2 := NewContext(ctx1)
+	defer cancel2()
 	if err := Run(ctx2); err != nil {
 		t.Fatal(err)
 	}
-
-	{
-		infos, err := Targets(ctx2)
-		if err != nil {
-			t.Fatal(err)
-		}
-		if want, got := 2, len(infos); want != got {
-			t.Fatalf("want %d targets, got %d", want, got)
-		}
-	}
+	wantTargets(ctx2, 2)
 
 	// The first context should also see both targets.
-	{
-		infos, err := Targets(ctx1)
-		if err != nil {
-			t.Fatal(err)
-		}
-		if want, got := 2, len(infos); want != got {
-			t.Fatalf("want %d targets, got %d", want, got)
-		}
-	}
+	wantTargets(ctx1, 2)
 }
diff --git a/target.go b/target.go
index b0ad051..225ea19 100644
--- a/target.go
+++ b/target.go
@@ -21,6 +21,7 @@ import (
 type Target struct {
 	browser   *Browser
 	SessionID target.SessionID
+	TargetID  target.ID
 
 	waitQueue  chan func(cur *cdp.Frame) bool
 	eventQueue chan *cdproto.Message