From 3bfec2211f5d8c180f5e83d758dd6245a7bc3521 Mon Sep 17 00:00:00 2001
From: Valentine <xb2w1z@gmail.com>
Date: Fri, 9 Aug 2024 14:21:56 +0300
Subject: [PATCH] fix error tests

---
 profile.go      | 22 ++++++++++++++++++----
 profile_test.go |  7 +++----
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/profile.go b/profile.go
index 631dfec..3fc0661 100644
--- a/profile.go
+++ b/profile.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"strings"
 	"sync"
 	"time"
 )
@@ -42,8 +43,9 @@ type user struct {
 	Data struct {
 		User struct {
 			Result struct {
-				RestID string     `json:"rest_id"`
-				Legacy legacyUser `json:"legacy"`
+				RestID  string     `json:"rest_id"`
+				Legacy  legacyUser `json:"legacy"`
+				Message string     `json:"message"`
 			} `json:"result"`
 		} `json:"user"`
 	} `json:"data"`
@@ -91,11 +93,17 @@ func (s *Scraper) GetProfile(username string) (Profile, error) {
 	}
 
 	if len(jsn.Errors) > 0 && jsn.Data.User.Result.RestID == "" {
+		if strings.Contains(jsn.Errors[0].Message, "Missing LdapGroup(visibility-custom-suspension)") {
+			return Profile{}, fmt.Errorf("user is suspended")
+		}
 		return Profile{}, fmt.Errorf("%s", jsn.Errors[0].Message)
 	}
 
 	if jsn.Data.User.Result.RestID == "" {
-		return Profile{}, fmt.Errorf("rest_id not found")
+		if jsn.Data.User.Result.Message == "User is suspended" {
+			return Profile{}, fmt.Errorf("user is suspended")
+		}
+		return Profile{}, fmt.Errorf("user not found")
 	}
 	jsn.Data.User.Result.Legacy.IDStr = jsn.Data.User.Result.RestID
 
@@ -142,11 +150,17 @@ func (s *Scraper) GetProfileByID(userID string) (Profile, error) {
 	}
 
 	if len(jsn.Errors) > 0 && jsn.Data.User.Result.RestID == "" {
+		if strings.Contains(jsn.Errors[0].Message, "Missing LdapGroup(visibility-custom-suspension)") {
+			return Profile{}, fmt.Errorf("user is suspended")
+		}
 		return Profile{}, fmt.Errorf("%s", jsn.Errors[0].Message)
 	}
 
 	if jsn.Data.User.Result.RestID == "" {
-		return Profile{}, fmt.Errorf("rest_id not found")
+		if jsn.Data.User.Result.Message == "User is suspended" {
+			return Profile{}, fmt.Errorf("user is suspended")
+		}
+		return Profile{}, fmt.Errorf("user not found")
 	}
 	jsn.Data.User.Result.Legacy.IDStr = jsn.Data.User.Result.RestID
 
diff --git a/profile_test.go b/profile_test.go
index db426d8..eb82891 100644
--- a/profile_test.go
+++ b/profile_test.go
@@ -1,7 +1,6 @@
 package twitterscraper_test
 
 import (
-	"fmt"
 	"strings"
 	"testing"
 	"time"
@@ -116,15 +115,15 @@ func TestGetProfileErrorSuspended(t *testing.T) {
 	if err == nil {
 		t.Error("Expected Error, got success")
 	} else {
-		if !strings.Contains(err.Error(), "Missing LdapGroup(visibility-custom-suspension)") {
-			t.Error("Expected error to contain 'Missing LdapGroup(visibility-custom-suspension)', got", err)
+		if !strings.Contains(err.Error(), "suspended") {
+			t.Error("Expected error to contain 'suspended', got", err)
 		}
 	}
 }
 
 func TestGetProfileErrorNotFound(t *testing.T) {
 	neUser := "sample3123131"
-	expectedError := fmt.Sprintf("User '%s' not found", neUser)
+	expectedError := "user not found"
 	_, err := testScraper.GetProfile(neUser)
 	if err == nil {
 		t.Error("Expected Error, got success")