[service-bus] Fix issue with live tests failing on session-based tests (

Azure#10325)

Two breakages:
* A service side change has made it so specifying different values for sessionId and partitionKey is now considered a failure.
* Fix a broken test that got through when I accidentally skipped the live test run on my BatchingReceiverLite PR

Also, moved the receiverInit.spec.ts into 'internal' since it is purely a unit test.

sdk/servicebus/service-bus/src/core/batchingReceiver.ts

@@ -44,7 +44,7 @@ export class BatchingReceiver extends MessageReceiver {
 
     this._batchingReceiverLite = new BatchingReceiverLite(
       context,
-      async (abortSignal?: AbortSignalLike): Promise<MinimalReceiver> => {
+      async (abortSignal?: AbortSignalLike): Promise<MinimalReceiver | undefined> => {
         let lastError: Error | AmqpError | undefined;
 
         const rcvrOptions = this._createReceiverOptions(false, {
@@ -67,7 +67,7 @@ export class BatchingReceiver extends MessageReceiver {
           throw lastError;
         }
 
-        return this._receiver!;
+        return this._receiver;
       },
       this.receiveMode
     );
@@ -241,7 +241,9 @@ interface ReceiveMessageArgs {
 export class BatchingReceiverLite {
   constructor(
     clientEntityContext: ClientEntityContext,
-    private _getCurrentReceiver: (abortSignal?: AbortSignalLike) => Promise<MinimalReceiver>,
+    private _getCurrentReceiver: (
+      abortSignal?: AbortSignalLike
+    ) => Promise<MinimalReceiver | undefined>,
     private _receiveMode: ReceiveMode
   ) {
     this._createServiceBusMessage = (context: MessageAndDelivery) => {
@@ -280,6 +282,12 @@ export class BatchingReceiverLite {
     try {
       this.isReceivingMessages = true;
       const receiver = await this._getCurrentReceiver(args.userAbortSignal);
+
+      if (receiver == null) {
+        // (was somehow closed in between the init() and the return)
+        return [];
+      }
+
       return await this._receiveMessagesImpl(receiver, args);
     } finally {
       this._closeHandler = undefined;

sdk/servicebus/service-bus/test/receiverInit.spec.ts → sdk/servicebus/service-bus/test/internal/receiverInit.spec.ts

@@ -6,9 +6,9 @@ import chaiAsPromised from "chai-as-promised";
 chai.use(chaiAsPromised);
 const assert = chai.assert;
 
-import { ClientEntityContext } from "../src/clientEntityContext";
-import { BatchingReceiver } from "../src/core/batchingReceiver";
-import { MessageReceiver, ReceiverType } from "../src/core/messageReceiver";
+import { ClientEntityContext } from "../../src/clientEntityContext";
+import { BatchingReceiver } from "../../src/core/batchingReceiver";
+import { MessageReceiver, ReceiverType } from "../../src/core/messageReceiver";
 
 describe("init() and close() interactions", () => {
   function fakeContext(): ClientEntityContext {

sdk/servicebus/service-bus/test/renewLockSessions.spec.ts

@@ -393,6 +393,7 @@ describe("Session Lock Renewal", () => {
   function getTestMessage(): ServiceBusMessage {
     const baseMessage = TestMessage.getSessionSample();
     baseMessage.sessionId = sessionId;
+    baseMessage.partitionKey = sessionId;
     return baseMessage;
   }
 });

sdk/servicebus/service-bus/test/streamingReceiverSessions.spec.ts

@@ -756,7 +756,7 @@ describe("Streaming with sessions", () => {
           body: "test",
           label: `${num}`,
           sessionId: TestMessage.sessionId,
-          partitionKey: "dummy" // Ensures all messages go to same partition to make peek work reliably
+          partitionKey: TestMessage.sessionId // Ensures all messages go to same partition to make peek work reliably
         };
         num++;
         messages.push(message);

sdk/servicebus/service-bus/test/utils/testUtils.ts

@@ -35,7 +35,7 @@ export class TestMessage {
     return {
       body: `message body ${randomNumber}`,
       messageId: `message id ${randomNumber}`,
-      partitionKey: `partition key ${randomNumber}`,
+      partitionKey: TestMessage.sessionId,
       contentType: `content type ${randomNumber}`,
       correlationId: `correlation id ${randomNumber}`,
       timeToLive: 60 * 60 * 24,