From 5b6ef4f1f2a51810386bea55278a785742133c96 Mon Sep 17 00:00:00 2001
From: Mattias Wadman <mattias.wadman@gmail.com>
Date: Fri, 1 Nov 2019 00:01:13 +0100
Subject: [PATCH] Make response type an/single/playlist

Make it possible to assert what response to get
---
 cmd/goutubedl/main.go |  4 +++-
 goutubedl.go          | 33 +++++++++++++++++++++++++++++----
 goutubedl_test.go     | 38 ++++++++++++++++++++++----------------
 3 files changed, 54 insertions(+), 21 deletions(-)

diff --git a/cmd/goutubedl/main.go b/cmd/goutubedl/main.go
index 924bd64..8a4c23b 100644
--- a/cmd/goutubedl/main.go
+++ b/cmd/goutubedl/main.go
@@ -12,12 +12,14 @@ import (
 )
 
 var dumpFlag = flag.Bool("J", false, "Dump JSON")
+var typeFlag = flag.String("t", "any", "Type")
 
 func main() {
 	log.SetFlags(0)
 	flag.Parse()
 
-	result, err := goutubedl.New(context.Background(), flag.Arg(0), goutubedl.Options{})
+	optType, _ := goutubedl.TypeFromString[*typeFlag]
+	result, err := goutubedl.New(context.Background(), flag.Arg(0), goutubedl.Options{Type: optType})
 	if err != nil {
 		log.Fatal(err)
 	}
diff --git a/goutubedl.go b/goutubedl.go
index bca1bfc..27aaa70 100644
--- a/goutubedl.go
+++ b/goutubedl.go
@@ -173,9 +173,27 @@ func (f Format) String() string {
 	)
 }
 
+// Type of response you want
+type Type int
+
+const (
+	// TypeAny single or playlist (default)
+	TypeAny Type = iota
+	// TypeSingle single track, file etc
+	TypeSingle
+	// TypePlaylist playlist with multiple tracks, files etc
+	TypePlaylist
+)
+
+var TypeFromString = map[string]Type{
+	"any":      TypeAny,
+	"single":   TypeSingle,
+	"playlist": TypePlaylist,
+}
+
 // Options for New()
 type Options struct {
-	YesPlaylist       bool // --yes-playlist
+	Type              Type
 	PlaylistStart     uint // --playlist-start
 	PlaylistEnd       uint // --playlist-end
 	DownloadThumbnail bool
@@ -234,7 +252,7 @@ func infoFromURL(ctx context.Context, rawURL string, options Options) (info Info
 		"--batch-file", "-",
 		"-J",
 	)
-	if options.YesPlaylist {
+	if options.Type == TypePlaylist {
 		cmd.Args = append(cmd.Args, "--yes-playlist")
 
 		if options.PlaylistStart > 0 {
@@ -300,8 +318,15 @@ func infoFromURL(ctx context.Context, rawURL string, options Options) (info Info
 		return Info{}, nil, infoErr
 	}
 
-	if options.YesPlaylist && (info.Type != "playlist" || info.Type == "multi_video") {
+	isPlaylist := info.Type == "playlist" || info.Type == "multi_video"
+	switch {
+	case options.Type == TypePlaylist && !isPlaylist:
 		return Info{}, nil, fmt.Errorf("not a playlist")
+	case options.Type == TypeSingle && isPlaylist:
+		return Info{}, nil, fmt.Errorf("not a single")
+	default:
+		// any type
+	}
 	}
 
 	// TODO: use headers from youtube-dl info for thumbnail and subtitle download?
@@ -334,7 +359,7 @@ func infoFromURL(ctx context.Context, rawURL string, options Options) (info Info
 	}
 
 	// as we ignore errors some entries might show up as null
-	if options.YesPlaylist {
+	if options.Type == TypePlaylist {
 		var filteredEntrise []Info
 		for _, e := range info.Entries {
 			if e.ID == "" {
diff --git a/goutubedl_test.go b/goutubedl_test.go
index 12bc708..63dd1bc 100644
--- a/goutubedl_test.go
+++ b/goutubedl_test.go
@@ -157,7 +157,7 @@ func TestPlaylist(t *testing.T) {
 	defer leakChecks(t)()
 
 	ydlResult, ydlResultErr := New(context.Background(), playlistRawURL, Options{
-		YesPlaylist:       true,
+		Type:              TypePlaylist,
 		DownloadThumbnail: false,
 	})
 
@@ -186,7 +186,7 @@ func TestPlaylistWithPrivateVideo(t *testing.T) {
 
 	playlistRawURL := "https://www.youtube.com/playlist?list=PLX0g748fkegS54oiDN4AXKl7BR7mLIydP"
 	ydlResult, ydlResultErr := New(context.Background(), playlistRawURL, Options{
-		YesPlaylist:       true,
+		Type:              TypePlaylist,
 		DownloadThumbnail: false,
 	})
 
@@ -201,20 +201,6 @@ func TestPlaylistWithPrivateVideo(t *testing.T) {
 	}
 }
 
-func TestPlaylistBadURL(t *testing.T) {
-	defer leakChecks(t)()
-
-	// using a non-playlist url
-	_, ydlResultErr := New(context.Background(), testVideoRawURL, Options{
-		YesPlaylist:       true,
-		DownloadThumbnail: false,
-	})
-
-	if ydlResultErr == nil {
-		t.Error("expected error")
-	}
-}
-
 func TestSubtitles(t *testing.T) {
 	defer leakChecks(t)()
 
@@ -246,3 +232,23 @@ func TestSubtitles(t *testing.T) {
 		}
 	}
 }
+
+func TestErrorNotAPlaylist(t *testing.T) {
+	_, ydlResultErr := New(context.Background(), testVideoRawURL, Options{
+		Type:              TypePlaylist,
+		DownloadThumbnail: false,
+	})
+	if ydlResultErr.Error() != "not a playlist" {
+		t.Errorf("expected is playlist error")
+	}
+}
+
+func TestErrorNotASingle(t *testing.T) {
+	_, ydlResultErr := New(context.Background(), playlistRawURL, Options{
+		Type:              TypeSingle,
+		DownloadThumbnail: false,
+	})
+	if ydlResultErr.Error() != "not a single" {
+		t.Errorf("expected is playlist error")
+	}
+}